Updated docs.

changes.md

@@ -6,6 +6,9 @@ nav_order: 3
 
 # Change Log
 
+## 0.6.0
+- Refactoring of training utilities. Added example for training an M3GNet potential.
+
 ## 0.5.6
 - Minor internal refactoring of basis expansions into `_basis.py`. (@lbluque)
 

docs/changes.md

@@ -6,6 +6,9 @@ nav_order: 3
 
 # Change Log
 
+## 0.6.0
+- Refactoring of training utilities. Added example for training an M3GNet potential.
+
 ## 0.5.6
 - Minor internal refactoring of basis expansions into `_basis.py`. (@lbluque)
 

docs/matgl.apps.md

@@ -14,7 +14,7 @@ potentials parameterizing the potential energy surface (PES).
 Implementation of Interatomic Potentials.
 
 
-### _class_ matgl.apps.pes.Potential(model: nn.Module, data_mean: torch.tensor | None = None, data_std: torch.tensor | None = None, element_refs: np.ndarray | None = None, calc_forces: bool = True, calc_stresses: bool = True, calc_hessian: bool = False)
+### _class_ matgl.apps.pes.Potential(model: nn.Module, data_mean: torch.Tensor | None = None, data_std: torch.Tensor | None = None, element_refs: np.ndarray | None = None, calc_forces: bool = True, calc_stresses: bool = True, calc_hessian: bool = False)
 Bases: `Module`, [`IOMixIn`](matgl.utils.md#matgl.utils.io.IOMixIn)
 
 A class representing an interatomic potential.
@@ -47,7 +47,7 @@ Initialize Potential from a model and elemental references.
 
 
 
-#### forward(g: dgl.DGLGraph, state_attr: torch.tensor | None = None, l_g: dgl.DGLGraph | None = None)
+#### forward(g: dgl.DGLGraph, state_attr: torch.Tensor | None = None, l_g: dgl.DGLGraph | None = None)
 
 * **Parameters**
 
@@ -64,7 +64,7 @@ Initialize Potential from a model and elemental references.
 
 * **Returns**
 
-    torch.tensor
+    torch.Tensor
 
 
 

docs/matgl.ext.md

@@ -55,7 +55,7 @@ Get a DGL graph from an input Atoms.
 
 
 
-### _class_ matgl.ext.ase.M3GNetCalculator(potential: [Potential](matgl.apps.md#matgl.apps.pes.Potential), state_attr: tensor | None = None, stress_weight: float = 1.0, \*\*kwargs)
+### _class_ matgl.ext.ase.M3GNetCalculator(potential: [Potential](matgl.apps.md#matgl.apps.pes.Potential), state_attr: Tensor | None = None, stress_weight: float = 1.0, \*\*kwargs)
 Bases: `Calculator`
 
 M3GNet calculator for ASE.
@@ -105,7 +105,7 @@ Perform calculation for an input Atoms.
 Properties calculator can handle (energy, forces, …)
 
 
-### _class_ matgl.ext.ase.MolecularDynamics(atoms: Atoms, potential: [Potential](matgl.apps.md#matgl.apps.pes.Potential), state_attr: torch.tensor = None, ensemble: str = 'nvt', temperature: int = 300, timestep: float = 1.0, pressure: float = 6.324209121801212e-07, taut: float | None = None, taup: float | None = None, compressibility_au: float | None = None, trajectory: str | Trajectory | None = None, logfile: str | None = None, loginterval: int = 1, append_trajectory: bool = False)
+### _class_ matgl.ext.ase.MolecularDynamics(atoms: Atoms, potential: [Potential](matgl.apps.md#matgl.apps.pes.Potential), state_attr: torch.Tensor = None, ensemble: str = 'nvt', temperature: int = 300, timestep: float = 1.0, pressure: float = 6.324209121801212e-07, taut: float | None = None, taup: float | None = None, compressibility_au: float | None = None, trajectory: str | Trajectory | None = None, logfile: str | None = None, loginterval: int = 1, append_trajectory: bool = False)
 Bases: `object`
 
 Molecular dynamics class.
@@ -125,7 +125,7 @@ Init the MD simulation.
     * **atoms** – 
 
 
-    * **state_attr** (*torch.tensor*) – State attr.
+    * **state_attr** (*torch.Tensor*) – State attr.
 
 
     * **ensemble** (*str*) – choose from ‘nvt’ or ‘npt’. NPT is not tested,
@@ -185,7 +185,7 @@ Set new atoms to run MD.
 
 
 
-### _class_ matgl.ext.ase.Relaxer(potential: [Potential](matgl.apps.md#matgl.apps.pes.Potential) = None, state_attr: torch.tensor = None, optimizer: Optimizer | str = 'FIRE', relax_cell: bool = True, stress_weight: float = 0.01)
+### _class_ matgl.ext.ase.Relaxer(potential: [Potential](matgl.apps.md#matgl.apps.pes.Potential) = None, state_attr: torch.Tensor = None, optimizer: Optimizer | str = 'FIRE', relax_cell: bool = True, stress_weight: float = 0.01)
 Bases: `object`
 
 Relaxer is a class for structural relaxation.
@@ -200,7 +200,7 @@ Relaxer is a class for structural relaxation.
     * **distribution** (*that comes with M3GNet*) – 
 
 
-    * **state_attr** (*torch.tensor*) – State attr.
+    * **state_attr** (*torch.Tensor*) – State attr.
 
 
     * **optimizer** (*str** or **ase Optimizer*) – the optimization algorithm.

docs/matgl.graph.md

@@ -57,8 +57,8 @@ Args:
 edges: DGL graph edges
 
 Returns:
-cos_theta: torch.tensor
-phi: torch.tensor
+cos_theta: torch.Tensor
+phi: torch.Tensor
 triple_bond_lengths (torch.tensor):
 
 
@@ -140,10 +140,20 @@ Returns: True if file exists.
 
 
 #### load(filename: str = 'dgl_graph.bin', filename_line_graph: str = 'dgl_line_graph.bin', filename_state_attr: str = 'state_attr.pt')
-Load dgl graphs
-Args:
-:filename: Name of file storing dgl graphs
-:filename: Name of file storing state attrs.
+Load dgl graphs from files.
+
+
+* **Parameters**
+
+
+    * **filename** – Name of file storing dgl graphs
+
+
+    * **filename_line_graph** – Name of file storing dgl line graphs
+
+
+    * **filename_state_attr** – Name of file storing state attrs.
+
 
 
 #### process()

docs/matgl.layers.md

@@ -32,7 +32,7 @@ Init SoftExponential with alpha value.
 
 
 
-#### forward(x: tensor)
+#### forward(x: Tensor)
 Evaluate activation function given the input tensor x.
 
 
@@ -66,7 +66,7 @@ softplus function that is 0 at x=0, the implementation aims at avoiding overflow
 Initializes the SoftPlus2 class.
 
 
-#### forward(x: tensor)
+#### forward(x: Tensor)
 Evaluate activation function given the input tensor x.
 
 
@@ -99,50 +99,51 @@ Bases: `Module`
 
 Get total property offset for a system.
 
-### Args:
-
+Args:
 property_offset (np.array): a array of elemental property offset.
 
 
-#### fit(structs_or_graphs: list, element_list: tuple[str], properties: np.typing.NDArray)
+#### fit(graphs: list, properties: np.typing.NDArray)
 Fit the elemental reference values for the properties.
 
 
 * **Parameters**
 
 
-    * **structs_or_graphs** – pymatgen Structures or dgl graphs
-
-
-    * **element_list** (*tuple*) – a list of element types
+    * **graphs** – dgl graphs
 
 
     * **properties** (*np.ndarray*) – array of extensive properties
 
 
 
-#### forward(g: dgl.DGLGraph, state_attr: torch.tensor | None = None)
+#### forward(g: dgl.DGLGraph, state_attr: torch.Tensor | None = None)
 Get the total property offset for a system.
 
-Args:
-g: a batch of dgl graphs
-state_attr: state attributes
 
-Returns:
-offset_per_graph:
+* **Parameters**
 
+
+    * **g** – a batch of dgl graphs
 
-#### get_feature_matrix(structs_or_graphs: list, element_list: tuple[str])
-Get the number of atoms for different elements in the structure.
 
+    * **state_attr** – state attributes
 
-* **Parameters**
 
-
-    * **structs_or_graphs** (*list*) – a list of pymatgen Structure or dgl graph
+
+* **Returns**
+
+    offset_per_graph
+
+
+
+#### get_feature_matrix(graphs: list)
+Get the number of atoms for different elements in the structure.
 
 
-    * **element_list** – a dictionary containing element types in the training set
+* **Parameters**
+
+    **graphs** (*list*) – a list of dgl graph
 
 
 
@@ -228,7 +229,7 @@ vanishes at cutoff.
 * **Parameters**
 
 
-    * **r** – torch.tensor pytorch tensors
+    * **r** – torch.Tensor pytorch tensors
 
 
     * **cutoff** – float, the cutoff radius
@@ -305,7 +306,7 @@ Ref:
 * **Parameters**
 
 
-    * **r** – torch.tensor distance tensor
+    * **r** – torch.Tensor distance tensor
 
 
     * **cutoff** – float, cutoff radius
@@ -338,7 +339,7 @@ num_centers (int): Number of centers for gaussian expansion.
 width (float): width of gaussian function.
 
 
-#### forward(bond_dist: tensor)
+#### forward(bond_dist: Tensor)
 Forward.
 
 Args:
@@ -373,7 +374,7 @@ Implementation of Set2Set.
 
 
 
-#### forward(g: DGLGraph, feat: tensor)
+#### forward(g: DGLGraph, feat: Tensor)
 Defines the computation performed at every call.
 
 
@@ -421,7 +422,7 @@ An implementation of a Gated multi-layer perceptron.
 
 
 
-#### forward(inputs: tensor)
+#### forward(inputs: Tensor)
 Defines the computation performed at every call.
 
 Should be overridden by all subclasses.
@@ -993,7 +994,7 @@ num_targets: number of target properties.
 
 * **Returns**
 
-    torch.tensor.
+    torch.Tensor.
 
 
 
@@ -1051,7 +1052,7 @@ Init ThreeBodyInteractions.
 
 
 
-#### forward(graph: DGLGraph, line_graph: DGLGraph, three_basis: tensor, three_cutoff: float, node_feat: tensor, edge_feat: tensor)
+#### forward(graph: DGLGraph, line_graph: DGLGraph, three_basis: Tensor, three_cutoff: float, node_feat: Tensor, edge_feat: Tensor)
 Forward function for ThreeBodyInteractions.
 
 
@@ -1097,10 +1098,10 @@ For the spherical Harmonics function, the column is ordered by
 * **Parameters**
 
 
-    * **sbf** – torch.tensor spherical bessel function results
+    * **sbf** – torch.Tensor spherical bessel function results
 
 
-    * **shf** – torch.tensor spherical harmonics function results
+    * **shf** – torch.Tensor spherical harmonics function results
 
 
     * **max_n** – int, max number of n

docs/matgl.md

@@ -588,52 +588,52 @@ MatGL (Materials Graph Library) is a graph deep learning library for materials s
     * [matgl.utils.training module](matgl.utils.md#module-matgl.utils.training)
 
 
-        * [`ModelTrainer`](matgl.utils.md#matgl.utils.training.ModelTrainer)
+        * [`MatglLightningModuleMixin`](matgl.utils.md#matgl.utils.training.MatglLightningModuleMixin)
 
 
-            * [`ModelTrainer.forward()`](matgl.utils.md#matgl.utils.training.ModelTrainer.forward)
+            * [`MatglLightningModuleMixin.configure_optimizers()`](matgl.utils.md#matgl.utils.training.MatglLightningModuleMixin.configure_optimizers)
 
 
-            * [`ModelTrainer.loss_fn()`](matgl.utils.md#matgl.utils.training.ModelTrainer.loss_fn)
+            * [`MatglLightningModuleMixin.on_test_model_eval()`](matgl.utils.md#matgl.utils.training.MatglLightningModuleMixin.on_test_model_eval)
 
 
-            * [`ModelTrainer.step()`](matgl.utils.md#matgl.utils.training.ModelTrainer.step)
+            * [`MatglLightningModuleMixin.on_train_epoch_end()`](matgl.utils.md#matgl.utils.training.MatglLightningModuleMixin.on_train_epoch_end)
 
 
-        * [`PotentialTrainer`](matgl.utils.md#matgl.utils.training.PotentialTrainer)
+            * [`MatglLightningModuleMixin.predict_step()`](matgl.utils.md#matgl.utils.training.MatglLightningModuleMixin.predict_step)
 
 
-            * [`PotentialTrainer.forward()`](matgl.utils.md#matgl.utils.training.PotentialTrainer.forward)
+            * [`MatglLightningModuleMixin.test_step()`](matgl.utils.md#matgl.utils.training.MatglLightningModuleMixin.test_step)
 
 
-            * [`PotentialTrainer.loss_fn()`](matgl.utils.md#matgl.utils.training.PotentialTrainer.loss_fn)
+            * [`MatglLightningModuleMixin.training_step()`](matgl.utils.md#matgl.utils.training.MatglLightningModuleMixin.training_step)
 
 
-            * [`PotentialTrainer.step()`](matgl.utils.md#matgl.utils.training.PotentialTrainer.step)
+            * [`MatglLightningModuleMixin.validation_step()`](matgl.utils.md#matgl.utils.training.MatglLightningModuleMixin.validation_step)
 
 
-        * [`TrainerMixin`](matgl.utils.md#matgl.utils.training.TrainerMixin)
+        * [`ModelLightningModule`](matgl.utils.md#matgl.utils.training.ModelLightningModule)
 
 
-            * [`TrainerMixin.configure_optimizers()`](matgl.utils.md#matgl.utils.training.TrainerMixin.configure_optimizers)
+            * [`ModelLightningModule.forward()`](matgl.utils.md#matgl.utils.training.ModelLightningModule.forward)
 
 
-            * [`TrainerMixin.on_test_model_eval()`](matgl.utils.md#matgl.utils.training.TrainerMixin.on_test_model_eval)
+            * [`ModelLightningModule.loss_fn()`](matgl.utils.md#matgl.utils.training.ModelLightningModule.loss_fn)
 
 
-            * [`TrainerMixin.on_train_epoch_end()`](matgl.utils.md#matgl.utils.training.TrainerMixin.on_train_epoch_end)
+            * [`ModelLightningModule.step()`](matgl.utils.md#matgl.utils.training.ModelLightningModule.step)
 
 
-            * [`TrainerMixin.predict_step()`](matgl.utils.md#matgl.utils.training.TrainerMixin.predict_step)
+        * [`PotentialLightningModule`](matgl.utils.md#matgl.utils.training.PotentialLightningModule)
 
 
-            * [`TrainerMixin.test_step()`](matgl.utils.md#matgl.utils.training.TrainerMixin.test_step)
+            * [`PotentialLightningModule.forward()`](matgl.utils.md#matgl.utils.training.PotentialLightningModule.forward)
 
 
-            * [`TrainerMixin.training_step()`](matgl.utils.md#matgl.utils.training.TrainerMixin.training_step)
+            * [`PotentialLightningModule.loss_fn()`](matgl.utils.md#matgl.utils.training.PotentialLightningModule.loss_fn)
 
 
-            * [`TrainerMixin.validation_step()`](matgl.utils.md#matgl.utils.training.TrainerMixin.validation_step)
+            * [`PotentialLightningModule.step()`](matgl.utils.md#matgl.utils.training.PotentialLightningModule.step)
 
 
         * [`xavier_init()`](matgl.utils.md#matgl.utils.training.xavier_init)

docs/matgl.models.md

@@ -116,7 +116,7 @@ The main M3GNet model.
 
 
 
-#### forward(g: dgl.DGLGraph, state_attr: torch.tensor | None = None, l_g: dgl.DGLGraph | None = None)
+#### forward(g: dgl.DGLGraph, state_attr: torch.Tensor | None = None, l_g: dgl.DGLGraph | None = None)
 Performs message passing and updates node representations.
 
 
@@ -145,7 +145,7 @@ Performs message passing and updates node representations.
 
 
 
-#### predict_structure(structure, state_feats: torch.tensor | None = None, graph_converter: [GraphConverter](matgl.graph.md#matgl.graph.converters.GraphConverter) | None = None)
+#### predict_structure(structure, state_feats: torch.Tensor | None = None, graph_converter: [GraphConverter](matgl.graph.md#matgl.graph.converters.GraphConverter) | None = None)
 Convenience method to directly predict property from structure.
 
 
@@ -297,7 +297,7 @@ Forward pass of MEGnet. Executes all blocks.
 
 
 
-#### predict_structure(structure, state_feats: torch.tensor | None = None, graph_converter: [GraphConverter](matgl.graph.md#matgl.graph.converters.GraphConverter) | None = None)
+#### predict_structure(structure, state_feats: torch.Tensor | None = None, graph_converter: [GraphConverter](matgl.graph.md#matgl.graph.converters.GraphConverter) | None = None)
 Convenience method to directly predict property from structure.