refactor: restructure dice environment handling and update Python version requirement

diceplayer/dice/dice_handler.py

@@ -1,14 +1,21 @@
+import warnings
+
 from diceplayer.dice.dice_input import (
     NPTEqConfig,
     NPTTerConfig,
     NVTEqConfig,
     NVTTerConfig,
 )
-from diceplayer.dice.dice_wrapper import DiceWrapper, DiceEnvironment
+from diceplayer.dice.dice_wrapper import (
+    DiceEnvironment,
+    DiceWrapper,
+)
+from diceplayer.environment import Atom, Molecule
 from diceplayer.logger import logger
 from diceplayer.state.state_model import StateModel
 
 import shutil
+from itertools import batched, chain, islice
 from pathlib import Path
 from threading import Thread
 
@@ -17,7 +24,7 @@ class DiceHandler:
     def __init__(self, step_directory: Path):
         self.dice_directory = step_directory / "dice"
 
-    def run(self, state: StateModel, cycle: int) -> list[DiceEnvironment]:
+    def run(self, state: StateModel, cycle: int) -> list[Atom]:
         if self.dice_directory.exists():
             logger.info(
                 f"Found dice directory: {self.dice_directory}, this directory will be purged for a clean state"
@@ -27,8 +34,8 @@ class DiceHandler:
 
         return self.run_simulations(state, cycle)
 
-    def run_simulations(self, state: StateModel, cycle: int) -> list[DiceEnvironment]:
-        results = []
+    def run_simulations(self, state: StateModel, cycle: int) -> list[Atom]:
+        results: list[list[Atom]] = []
 
         threads = []
         for p in range(state.config.dice.nprocs):
@@ -44,12 +51,14 @@ class DiceHandler:
                 f"Expected {state.config.dice.nprocs} simulation results, but got {len(results)}"
             )
 
-        return [
-            i for i in [r for r in results]
-        ]
+        return self._aggregate_results(state, results)
 
     def _simulation_process(
-        self, state: StateModel, cycle: int, proc: int, results: list[list[DiceEnvironment]]
+        self,
+        state: StateModel,
+        cycle: int,
+        proc: int,
+        results: list[list[Atom]],
     ) -> None:
         proc_directory = self.dice_directory / f"{proc:02d}"
         if proc_directory.exists():
@@ -77,7 +86,12 @@ class DiceHandler:
             npt_eq_config = NPTEqConfig.from_config(state.config)
             dice.run(npt_eq_config)
 
-        results.append(dice.parse_results(state.system))
+        results.extend(
+            [
+                self._filter_environment_sites(state, environment)
+                for environment in dice.parse_results()
+            ]
+        )
 
     @staticmethod
     def _generate_phb_file(state: StateModel, proc_directory: Path) -> None:
@@ -106,3 +120,67 @@ class DiceHandler:
                     )
 
     def _generate_last_xyz(self, state: StateModel, proc_directory: Path) -> None: ...
+
+    @staticmethod
+    def _filter_environment_sites(
+        state: StateModel, environment: DiceEnvironment
+    ) -> list[Atom]:
+        picked_environment = []
+
+        ref_molecule = state.system.molecule[0]
+        ref_molecule_sizes = ref_molecule.sizes_of_molecule()
+        ref_n_sites = len(ref_molecule.atom) * state.config.dice.nmol[0]
+
+        min_distance = min(
+            (environment.thickness[i] - ref_molecule_sizes[i]) / 2 for i in range(3)
+        )
+
+        site_iter = iter(environment.items)
+        _ = list(islice(site_iter, ref_n_sites))
+
+        for molecule_index, molecule in enumerate(state.system.molecule[1:], start=1):
+            molecule_n_atoms = len(molecule.atom)
+            molecule_n_sites = molecule_n_atoms * state.config.dice.nmol[molecule_index]
+
+            sites = list(islice(site_iter, molecule_n_sites))
+
+            for molecule_sites in batched(sites, molecule_n_atoms):
+                new_molecule = Molecule("ASEC TMP MOLECULE")
+
+                for site_index, atom_site in enumerate(molecule_sites):
+                    new_molecule.add_atom(
+                        Atom(
+                            molecule.atom[site_index].lbl,
+                            molecule.atom[site_index].na,
+                            atom_site.x,
+                            atom_site.y,
+                            atom_site.z,
+                            molecule.atom[site_index].chg,
+                            molecule.atom[site_index].eps,
+                            molecule.atom[site_index].sig,
+                        )
+                    )
+
+                if molecule.signature() != new_molecule.signature():
+                    _message = f"Skipping sites because the molecule signature does not match the reference molecule. Expected {molecule.signature()} but got {new_molecule.signature()}"
+                    warnings.warn(_message)
+                    logger.warning(_message)
+                    continue
+
+                if ref_molecule.minimum_distance(new_molecule) >= min_distance:
+                    continue
+
+                picked_environment.extend(new_molecule.atom)
+
+        return picked_environment
+
+    @staticmethod
+    def _aggregate_results(state: StateModel, results: list[list[Atom]]) -> list[Atom]:
+        norm_factor = round(state.config.dice.nstep[-1] / state.config.dice.isave)
+
+        agg_results = []
+        for atom in chain(*[r for r in results]):
+            atom.chg = atom.chg * norm_factor
+            agg_results.append(atom)
+
+        return agg_results