Adjust code and fix script saving after changes

Author: MBartkowiakSTFC

MDANSE/Src/MDANSE/Chemistry/ChemicalSystem.py

@@ -19,11 +19,11 @@
 import itertools as it
 from collections import defaultdict
 from collections.abc import Iterable
-from functools import reduce
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, SupportsInt
 
 import h5py
+import more_itertools
 import networkx as nx
 import numpy as np
 from more_itertools import padded
@@ -439,7 +439,7 @@ def find_clusters_from_bonds(self):
         while len(atom_pool) > 0:
             last_atom = atom_pool.pop()
             temp_dict = nx.dfs_successors(total_graph, last_atom)
-            others = reduce(list.__iadd__, temp_dict.values(), [])
+            others = list(more_itertools.flatten(temp_dict.values()))
             for atom in others:
                 atom_pool.remove(atom)
             molecule = [last_atom, *others]

MDANSE/Src/MDANSE/Framework/Configurators/HDFTrajectoryConfigurator.py

@@ -15,6 +15,7 @@
 #
 from __future__ import annotations
 
+import bisect
 from pathlib import Path
 
 import h5py
@@ -23,7 +24,7 @@
 from MDANSE import PLATFORM
 from MDANSE.Framework.Configurators.IConfigurator import IConfigurator
 from MDANSE.Framework.Configurators.InputFileConfigurator import InputFileConfigurator
-from MDANSE.MolecularDynamics.Trajectory import Trajectory, check_hdf5_driver
+from MDANSE.MolecularDynamics.Trajectory import Trajectory
 
 TIME_STEP_TOL = 1e-8
 DATASET_CACHE_SIZE = 2**24
@@ -98,10 +99,8 @@ def next_prime(nchunks: int) -> int:
     int
         A prime number larger than the input number.
     """
-    import bisect
-
-    ip = bisect.bisect_right(nchunks, PRIMES)
-    return nchunks if ip == len(PRIMES) else PRIMES[ip]
+    ip = bisect.bisect_right(PRIMES, nchunks)
+    return nchunks if ip == len(PRIMES) else PRIMES[ip]
 
 
 def guess_hdf5_trajectory_parameters(
@@ -160,12 +159,12 @@ def configure(self, value):
                 driver = None
                 rdcc_nbytes, rdcc_nslots = guess_hdf5_trajectory_parameters(value)
                 rdcc_w0 = None
-            case (str(), str(), int(), int(), int()):
+            case (str() | Path(), str(), int(), int(), float()):
                 file_name, driver, rdcc_nbytes, rdcc_nslots, rdcc_w0 = value
                 driver = driver if driver in HDF5_DRIVERS else None
             case _:
                 self.error_status = f"Invalid value {value!r}"
-                return 
+                return
 
         self["driver"] = driver
         self["rdcc_nbytes"] = rdcc_nbytes

MDANSE/Src/MDANSE/Framework/Jobs/TrajectoryEditor.py

@@ -15,8 +15,7 @@
 #
 from __future__ import annotations
 
-from functools import reduce
-
+import more_itertools
 import numpy as np
 
 from MDANSE.Chemistry.ChemicalSystem import (
@@ -160,16 +159,17 @@ def initialize(self):
                 com_conf = AbsoluteConfiguration(
                     coords,
                 )
-            self.grouped_indices = list(more_itertools.flatten(new_chemical_system.clusters.values()))
+            self.grouped_indices = list(
+                more_itertools.flatten(new_chemical_system.clusters.values())
+            )
             coords = com_conf.contiguous_configuration(self.grouped_indices).coordinates
         else:
             assign_molecules_after_atom_selection(
                 self._indices, self._input_chemical_system, new_chemical_system
             )
-            self.grouped_indices = reduce(
-                list.__add__, new_chemical_system.clusters.values(), []
+            self.grouped_indices = list(
+                more_itertools.flatten(new_chemical_system.clusters.values())
             )
-
         # The output trajectory is opened for writing.
         self._output_trajectory = TrajectoryWriter(
             self.configuration["output_files"]["file"],

MDANSE/Src/MDANSE/MolecularDynamics/Configuration.py

@@ -18,11 +18,12 @@
 import abc
 import copy
 from collections.abc import Sequence
-from functools import reduce
 from typing import TYPE_CHECKING, Any
 
+import more_itertools
 import networkx as nx
 import numpy as np
+import typing_extensions
 
 from MDANSE.MLogging import LOG
 
@@ -199,7 +200,7 @@ def continuous_coordinates(
     while len(atom_pool) > 0:
         last_atom = atom_pool.pop()
         temp_dict = nx.dfs_successors(total_graph, last_atom)
-        others = reduce(list.__add__, temp_dict.values(), [])
+        others = list(more_itertools.flatten(temp_dict.values()))
         for atom in others:
             atom_pool.pop(atom_pool.index(atom))
         segment = [last_atom, *others]
@@ -291,7 +292,6 @@ class _Configuration(metaclass=abc.ABCMeta):
     is_periodic: bool
 
     def __init__(self, coords: ArrayLike, **variables):
-
         self._variables = {}
 
         self["coordinates"] = np.array(coords, dtype=float)
@@ -416,7 +416,7 @@ def __init__(
             raise ValueError("Invalid unit cell dimensions")
         self._unit_cell = unit_cell
 
-    def clone(self) -> _PeriodicConfiguration:
+    def clone(self) -> typing_extensions.Self:
         """Return a deep copy of this configuration."""
 
         unit_cell = copy.deepcopy(self._unit_cell)

MDANSE/Src/MDANSE/MolecularDynamics/Trajectory.py

@@ -17,7 +17,6 @@
 
 import copy
 import html
-import math
 import os
 from collections import Counter, defaultdict
 from enum import auto
@@ -177,9 +176,10 @@ def __init__(self,
                  filename,
                  trajectory_format: ValidFormats | None = None,
                  hdf5_driver: str | None = None,
+                 *,
                  rdcc_nbytes: int | None = None,
-                 rdcc_w0: float | None = None,
                  rdcc_nslots: int | None = None,
+                 rdcc_w0: float | None = None,
                  fast_load: bool = False):
         self._filename = filename
         self._hdf5_driver = hdf5_driver
@@ -192,9 +192,9 @@ def __init__(self,
 
         self._trajectory = self.open_trajectory(self._format,
                                                     self._hdf5_driver,
-                                                    self._rdcc_nbytes,
-                                                    self._rdcc_w0,
-                                                    self._rdcc_nslots,
+                                                    rdcc_nbytes=self._rdcc_nbytes,
+                                                    rdcc_nslots=self._rdcc_nslots,
+                                                    rdcc_w0=self._rdcc_w0,
                                                     fast_load = fast_load
                                                     )
         self._min_span = None
@@ -469,9 +469,10 @@ def guess_correct_format(self) -> ValidFormats:
     def open_trajectory(self,
                         trajectory_format,
                         hdf5_driver,
-                        rdcc_nbytes,
-                        rdcc_w0,
-                        rdcc_nslots,
+                        *,
+                        rdcc_nbytes: int | None = None,
+                        rdcc_nslots: int | None = None,
+                        rdcc_w0: float | None = None,
                         fast_load: bool = False
                         ):
         trajectory_class = available_formats[trajectory_format]
@@ -510,7 +511,12 @@ def __getstate__(self):
 
     def __setstate__(self, state):
         self.__dict__ = state
-        self._trajectory = self.open_trajectory(self._format, self._hdf5_driver, self._rdcc_nbytes, self._rdcc_w0, self._rdcc_nslots, fast_load=True)
+        self._trajectory = self.open_trajectory(self._format,
+                                                self._hdf5_driver,
+                                                rdcc_nbytes=self._rdcc_nbytes,
+                                                rdcc_nslots=self._rdcc_nslots,
+                                                rdcc_w0=self._rdcc_w0,
+                                                fast_load=True)
 
     def __len__(self):
         return len(self._trajectory)
@@ -1026,10 +1032,12 @@ def __init__(
         chemical_system: ChemicalSystem,
         n_steps,
         selected_atoms=None,
+        *,
         positions_dtype=np.float64,
         chunking_limit=(1,128),
         compression="none",
         initial_charges=None,
+        meta_block_size: int = 65536
     ):
         """Constructor.
 
@@ -1044,7 +1052,7 @@ def __init__(
         """
         self._h5_filename = Path(h5_filename)
         PLATFORM.create_directory(self._h5_filename.parent)
-        self._h5_file = h5py.File(self._h5_filename, "w", meta_block_size=65536)
+        self._h5_file = h5py.File(self._h5_filename, "w", meta_block_size=meta_block_size)
 
         self._chemical_system = chemical_system
         self._last_configuration = None

MDANSE/Src/MDANSE/MolecularDynamics/TrajectoryUtils.py

@@ -15,12 +15,11 @@
 #
 from __future__ import annotations
 
-from functools import reduce
 from itertools import pairwise
 from typing import TYPE_CHECKING
 
 import numpy as np
-from more_itertools import chunked_even
+from more_itertools import chunked_even, flatten
 
 if TYPE_CHECKING:
     from MDANSE.Chemistry.ChemicalSystem import ChemicalSystem
@@ -130,7 +129,7 @@ def group_cluster_indices(chemical_system: ChemicalSystem) -> list[list[int]]:
     list[list[int]]
         A list of atom index lists, one per molecule instance.
     """
-    return reduce(list.__add__, chemical_system.clusters.values(), [])
+    return list(flatten(chemical_system.clusters.values()))
 
 
 def group_atom_indices(

MDANSE/Src/MDANSE/Trajectory/H5MDTrajectory.py

@@ -17,11 +17,11 @@
 
 from collections import ChainMap
 from enum import Enum
-from functools import reduce
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 import h5py
+import more_itertools
 import numpy as np
 import numpy.typing as npt
 
@@ -135,9 +135,10 @@ def __init__(
         self,
         h5_filename: Path | str,
         hdf5_driver: str | None = None,
+        *,
         rdcc_nbytes: int | None = None,
-        rdcc_w0: float | None = None,
         rdcc_nslots: int | None = None,
+        rdcc_w0: float | None = None,
         fast_load: bool = False,
     ):
         """Constructor.
@@ -158,8 +159,8 @@ def __init__(
             "r",
             driver=hdf5_driver,
             rdcc_nbytes=rdcc_nbytes,
-            rdcc_w0=rdcc_w0,
             rdcc_nslots=rdcc_nslots,
+            rdcc_w0=rdcc_w0,
         )
 
         particle_types = self._h5_file["/particles/all/species"]
@@ -204,8 +205,8 @@ def __init__(
 
         if self._chemical_system.rdkit_mol.GetNumBonds() > 0:
             configuration = self.configuration(0)
-            grouped_indices = reduce(
-                list.__add__, self._chemical_system.clusters.values(), []
+            grouped_indices = list(
+                more_itertools.flatten(self._chemical_system.clusters.values())
             )
             contiguous_configuration = configuration.contiguous_configuration(
                 grouped_indices

MDANSE/Src/MDANSE/Trajectory/MdanseTrajectory.py

@@ -15,15 +15,15 @@
 #
 from __future__ import annotations
 
-from collections import ChainMap, defaultdict
-from collections.abc import Iterable, Mapping
-from functools import cached_property, reduce
+from collections import ChainMap
+from collections.abc import Mapping
+from functools import cached_property
 from pathlib import Path
 
 import h5py
 import numpy as np
 import numpy.typing as npt
-from more_itertools import first
+from more_itertools import first, flatten
 
 import MDANSE
 from MDANSE.Chemistry import ATOMS_DATABASE
@@ -65,9 +65,10 @@ def __init__(
         self,
         h5_filename: Path | str,
         hdf5_driver: str | None = None,
+        *,
         rdcc_nbytes: int | None = None,
-        rdcc_w0: float | None = None,
         rdcc_nslots: int | None = None,
+        rdcc_w0: float | None = None,
         fast_load: bool = False,
     ):
         """Open the file and build a trajectory.
@@ -94,8 +95,8 @@ def __init__(
             "r",
             driver=hdf5_driver,
             rdcc_nbytes=rdcc_nbytes,
-            rdcc_w0=rdcc_w0,
             rdcc_nslots=rdcc_nslots,
+            rdcc_w0=rdcc_w0,
         )
         self._has_database = "atom_database" in self._h5_file
         self._has_atoms = []
@@ -113,9 +114,7 @@ def __init__(
 
         if self._chemical_system.rdkit_mol.GetNumBonds() > 0:
             configuration = self.configuration(0)
-            grouped_indices = reduce(
-                list.__add__, self._chemical_system.clusters.values(), []
-            )
+            grouped_indices = list(flatten(self._chemical_system.clusters.values()))
             contiguous_configuration = configuration.contiguous_configuration(
                 grouped_indices
             )

MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/HDFTrajectoryWidget.py

@@ -162,12 +162,11 @@ def get_widget_value(self):
         else:
             self._label.setToolTip(self._tooltip)
         hdf5_driver = self._driver_widget.currentText()
-        hdf5_driver = None if hdf5_driver == "default" else hdf5_driver
         rdcc = {}
         for name, widget in self._extra_widgets.items():
             rdcc[name] = widget.value() if widget.isEnabled() else None
         return (
-            result,
+            str(result),
             hdf5_driver,
             rdcc["rdcc_nbytes"] * 1024**2 if rdcc["rdcc_nbytes"] is not None else None,
             rdcc["rdcc_nslots"],

MDANSE_GUI/Src/MDANSE_GUI/InputWidgets/OutputTrajectoryWidget.py

@@ -69,8 +69,9 @@ def __init__(self, *args, **kwargs):
         self.chunk_atom_box = QSpinBox(self._base)
         self.chunk_frame_box = QSpinBox(self._base)
         for typ, chunk_box in zip(
-           ("atoms", "frames"),
-           (self.chunk_atom_box, self.chunk_frame_box)
+            ("atoms", "frames"),
+            (self.chunk_atom_box, self.chunk_frame_box),
+            strict=True,
         ):
             chunk_box.setMinimum(1)
             chunk_box.setMaximum(0xFFFF)