Archive data update

simulationworkflowschema/geometry_optimization.py

@@ -17,17 +17,16 @@
 #
 import numpy as np
 from ase import Atoms
-
-from nomad.metainfo import SubSection, Quantity, MEnum
 from nomad.datamodel.metainfo.workflow import Link
+from nomad.metainfo import MEnum, Quantity, SubSection
 from runschema.calculation import EnergyEntry
 
 from .general import (
-    SimulationWorkflowMethod,
-    SimulationWorkflowResults,
-    SerialSimulation,
     WORKFLOW_METHOD_NAME,
     WORKFLOW_RESULTS_NAME,
+    SerialSimulation,
+    SimulationWorkflowMethod,
+    SimulationWorkflowResults,
     resolve_difference,
 )
 
@@ -61,7 +60,7 @@ class GeometryOptimizationMethod(SimulationWorkflowMethod):
         shape=[],
         description="""
         The method used for geometry optimization. Some known possible values are:
-        `"steepest_descent"`, `"conjugant_gradient"`, `"low_memory_broyden_fletcher_goldfarb_shanno"`.
+        `"steepest_descent"`, `"conjugate_gradient"`, `"low_memory_broyden_fletcher_goldfarb_shanno"`.
         """,
     )
 

simulationworkflowschema/molecular_dynamics.py

@@ -15,50 +15,52 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Any, Callable, Optional
-from itertools import chain
-from collections import namedtuple
-import numpy as np
 from array import array
-from scipy import sparse
-from scipy.stats import linregress
+from collections import namedtuple
+from itertools import chain
+from typing import Any, Callable, Optional
+
+import ase
+import MDAnalysis
+import MDAnalysis.analysis.rdf as MDA_RDF
 import networkx
 import numpy as np
-import MDAnalysis
+from MDAnalysis.core._get_readers import get_reader_for
 from MDAnalysis.core.topology import Topology
 from MDAnalysis.core.universe import Universe
-import MDAnalysis.analysis.rdf as MDA_RDF
-from MDAnalysis.core._get_readers import get_reader_for
-
+from nomad import atomutils
 from nomad.datamodel.data import ArchiveSection
+from nomad.datamodel.hdf5 import HDF5Dataset
+from nomad.datamodel.metainfo.workflow import Link
 from nomad.metainfo import (
-    SubSection,
-    Section,
-    Quantity,
     MEnum,
-    Reference,
     MSection,
+    Quantity,
+    Reference,
+    Section,
+    SubSection,
 )
-from nomad.datamodel.hdf5 import HDF5Dataset
-from nomad.datamodel.metainfo.workflow import Link
-from runschema.system import System, AtomsGroup
+from nomad.units import ureg
+from nomad.utils import get_logger
 from runschema.calculation import (
     RadiusOfGyration as RadiusOfGyrationCalculation,
+)
+from runschema.calculation import (
     RadiusOfGyrationValues as RadiusOfGyrationValuesCalculation,
 )
-from nomad.utils import get_logger
-from nomad.units import ureg
-from nomad import atomutils
+from runschema.system import AtomsGroup, System
+from scipy import sparse
+from scipy.stats import linregress
+
 from .general import (
-    SimulationWorkflowMethod,
-    SimulationWorkflowResults,
-    SerialSimulation,
     WORKFLOW_METHOD_NAME,
     WORKFLOW_RESULTS_NAME,
+    SerialSimulation,
+    SimulationWorkflowMethod,
+    SimulationWorkflowResults,
 )
 from .thermodynamics import ThermodynamicsResults
 
-
 LOGGER = get_logger(__name__)
 
 
@@ -305,36 +307,53 @@ def archive_to_universe(
 
         bonds (tuple, shape=([])): list of tuples with the atom indices of each bond
     """
-
-    try:
-        sec_run = archive.run[-1]
-        sec_system = sec_run.system
-        sec_system_top = sec_run.system[system_index]
-        sec_atoms = sec_system_top.atoms
-        sec_atoms_group = sec_system_top.atoms_group
-        sec_calculation = sec_run.calculation
-        sec_method = (
-            sec_run.method[method_index] if sec_run.get('method') is not None else {}
+    print(f'Nomad schema: {archive}')
+    if archive.data:
+        print('Found archive.data in nomad-schema')
+        data = archive.data
+        sec_system = data.model_system  # full list — iterated per frame
+        sec_system_top = sec_system[system_index]  # topology frame
+        sec_atoms_group = (
+            sec_system_top.sub_systems if sec_system_top is not None else None
         )
-    except IndexError:
+        sec_method = data.model_method[method_index] if data.model_method else None
+        sec_outputs = data.outputs if data.outputs else None
+    else:
         LOGGER.warning(
-            'Supplied indices or necessary sections do not exist in archive. Cannot build the MDA universe.'
+            'No data section found in archive. Cannot build the MDA universe.'
         )
         return None
 
-    n_atoms = sec_atoms.get('n_atoms')
-    if n_atoms is None:
-        LOGGER.warning('No atoms found in the archive. Cannot build the MDA universe.')
-        return None
-
-    n_frames = len(sec_system) if sec_system is not None else 1
-    atom_names = sec_atoms.get('labels')
-    model_atom_parameters = sec_method.get('atom_parameters')
+    n_atoms = sec_system_top.n_particles if sec_system_top is not None else None
+    particle_states = (
+        sec_system_top.particle_states if sec_system_top is not None else None
+    )
+    atom_names = (
+        [ps.label or 'CGX' for ps in particle_states] if particle_states else None
+    )
     atom_types = (
-        [atom.label for atom in model_atom_parameters]
-        if model_atom_parameters
-        else atom_names
+        [ps.chemical_symbol or 'CGX' for ps in particle_states]
+        if particle_states
+        else None
+    )
+    masses = [
+        ureg.convert(ps.mass.magnitude, ps.mass.units, ureg.amu)
+        if ps.mass is not None
+        else ase.data.atomic_masses[
+            ase.data.atomic_numbers.get(ps.chemical_symbol or '', 0)
+        ]
+        for ps in particle_states
+    ]
+    charges = [
+        ureg.convert(ps.partial_charge.magnitude, ps.partial_charge.units, ureg.e)
+        if ps.partial_charge is not None
+        else 0.0
+        for ps in particle_states
+    ]
+    system_times = (
+        [out.time for out in sec_outputs if out.time is not None] if sec_outputs else []
     )
+    n_frames = len(sec_system) if sec_system is not None else 1
     atom_resindex = np.arange(n_atoms)
     atoms_segindices = np.empty(n_atoms)
     atom_segids = np.array(range(n_atoms), dtype='object')
@@ -351,6 +370,9 @@ def archive_to_universe(
     molecule_n_res = []
     for mol_group_ind, mol_group in enumerate(molecule_groups):
         atoms_segindices[mol_group.atom_indices] = mol_group_ind
+        print(
+            f'mol_group.label: {mol_group.label}, mol_group.atom_indices: {mol_group.atom_indices}'
+        )
         atom_segids[mol_group.atom_indices] = mol_group.label
         molecules = mol_group.atoms_group if mol_group.atoms_group is not None else []
         for mol in molecules:
@@ -1867,7 +1889,7 @@ class MolecularDynamicsMethod(SimulationWorkflowMethod):
     integrator_type = Quantity(
         type=MEnum(
             'brownian',
-            'conjugant_gradient',
+            'conjugate_gradient',
             'langevin_goga',
             'langevin_schneider',
             'leap_frog',