Parallel Differential Equation runner

src/NanoParticleTools/builder.py

@@ -3,8 +3,7 @@
                                         get_spectrum_wavelength_from_dndt,
                                         average_dndt, intensities_from_docs)
 
-from maggma.core import Builder
-from maggma.core import Store
+from maggma.core import Builder, Store
 from maggma.utils import grouper
 from typing import Iterator, List, Dict, Optional, Iterable, Tuple
 from bson import uuid

src/NanoParticleTools/differential_kinetics/__init__.py

@@ -0,0 +1,4 @@
+from NanoParticleTools.differential_kinetics.util import (
+    get_templates, run_one_rate_eq, save_data_to_hdf5, load_data_from_hdf5,
+    run_and_save_one, get_diff_kinetics_parser)
+from NanoParticleTools.differential_kinetics.runner import DifferentialKinetics

src/NanoParticleTools/differential_kinetics/runner.py

@@ -0,0 +1,60 @@
+from NanoParticleTools.differential_kinetics import (get_templates,
+                                                     run_one_rate_eq,
+                                                     save_data_to_hdf5)
+from NanoParticleTools.species_data import Dopant
+from maggma.core import Builder, Store
+from argparse import ArgumentParser
+from h5py import File
+
+from typing import Iterator
+
+
+class DifferentialKinetics(Builder):
+    """
+    Builder that processes and averages NPMC documents
+    """
+
+    def __init__(self, args: ArgumentParser, **kwargs):
+
+        self.source = None
+        self.args = args
+        self.target = None
+        self.kwargs = kwargs
+        self._file = None
+
+        super().__init__(sources=None, targets=None, chunk_size=1, **kwargs)
+
+    def connect(self):
+        # Since we aren't using stores, do nothing
+        return
+
+    @property
+    def file(self):
+        if self._file is None:
+            self._file = File(self.args.output_file, 'w')
+        return self._file
+
+    def get_items(self) -> Iterator[dict]:
+        for sample_id, template in enumerate(get_templates(self.args)):
+            yield (sample_id, template)
+
+    def process_item(self, item: tuple[int, dict]) -> dict:
+        sample_id, template = item
+        dopants = [
+            Dopant(el, x)
+            for el, x in zip(template['dopants'], template['dopant_concs'])
+        ]
+        output = run_one_rate_eq(
+            dopants,
+            excitation_wavelength=template['excitation_wavelength'],
+            excitation_power=template['excitation_power'],
+            include_spectra=self.args.include_spectra)
+
+        group_id = int(sample_id // self.args.max_data_per_group)
+        data_id = int(sample_id % self.args.max_data_per_group)
+        return (group_id, data_id, output)
+
+    def update_targets(self, items: list[dict]) -> None:
+        for item in items:
+            group_id, data_id, output = item
+            save_data_to_hdf5(self.file, group_id, data_id, output)

src/NanoParticleTools/differential_kinetics.py → src/NanoParticleTools/differential_kinetics/util.py

@@ -14,8 +14,10 @@
 import json
 import h5py
 
+import logging
 
-def get_parser():
+
+def get_diff_kinetics_parser():
     parser = argparse.ArgumentParser()
 
     parser.add_argument('-n',
@@ -85,6 +87,35 @@ def get_parser():
     return parser
 
 
+def get_templates(args):
+    for sample_id in range(args.num_samples):
+        # Pick a number of dopants
+        n_dopants = np.random.choice(range(1, args.max_dopants + 1))
+
+        # Pick the dopants
+        dopants = np.random.choice(args.dopants, n_dopants, replace=False)
+
+        # Get the dopant concentrations, normalizing the total concentration to 0-1
+        total_conc = np.random.uniform(0, 1)
+        dopant_concs = np.random.uniform(0, 1, n_dopants)
+        dopant_concs = total_conc * dopant_concs / np.sum(dopant_concs)
+
+        # sample a wavelength
+        wavelength = np.random.uniform(*args.excitation_wavelength)
+
+        # sample a power
+        power = np.random.uniform(*args.excitation_power)
+        yield {
+            # 'sample_id': sample_id,
+            # 'group_id': sample_id // args.max_data_per_group,
+            # 'data_id': sample_id % args.max_data_per_group,
+            'dopants': dopants,
+            'dopant_concs': dopant_concs,
+            'excitation_wavelength': wavelength,
+            'excitation_power': power
+        }
+
+
 def run_one_rate_eq(dopants: list[Dopant],
                     excitation_wavelength: int,
                     excitation_power: float,
@@ -128,9 +159,10 @@ def run_one_rate_eq(dopants: list[Dopant],
             'simulation_time': t[-1],
             'excitation_wavelength': excitation_wavelength,
             'excitation_power': excitation_power,
-            'dopant_concentrations':
-            {dopant.symbol: dopant.molar_concentration
-             for dopant in dopants}
+            'dopant_concentrations': {
+                dopant.symbol: dopant.molar_concentration
+                for dopant in dopants
+            }
         }
     }
     out_dict['populations'] = final_pop
@@ -152,11 +184,14 @@ def run_one_rate_eq(dopants: list[Dopant],
 
 
 def run_and_save_one(dopants: list[str], dopant_concs: list[float],
-                     group_id: int, data_i: int, file: h5py.File, **kwargs):
+                     group_id: int, data_id: int, sample_id: int,
+                     file: h5py.File, **kwargs):
     dopants = [Dopant(el, x) for el, x in zip(dopants, dopant_concs)]
+
+    logging.info(f'Running Sample {sample_id}')
     out_dict = run_one_rate_eq(dopants, **kwargs)
 
-    save_data_to_hdf5(file, group_id, data_i, out_dict)
+    save_data_to_hdf5(file, group_id, data_id, out_dict)
 
     return out_dict
 

tests/differential_kinetics/test_runner.py

@@ -0,0 +1,19 @@
+from NanoParticleTools.differential_kinetics.runner import DifferentialKinetics
+from NanoParticleTools.differential_kinetics.util import (
+    get_diff_kinetics_parser, load_data_from_hdf5)
+from h5py import File
+
+
+def test_runner(tmp_path):
+    args = get_diff_kinetics_parser().parse_args([
+        '-n', '2', '-w', '400', '700', '-p', '10', '100000', '-o',
+        f'{tmp_path}/out.h5', '-s'
+    ])
+    dk = DifferentialKinetics(args)
+    dk.run()
+
+    data_points = []
+    with File(tmp_path / 'out.h5', 'r') as f:
+        for i in range(0, len(f['group_0'])):
+            data_points.append(load_data_from_hdf5(f, 0, i))
+    assert len(data_points) == 2

tests/test_differential_kinetics.py → tests/differential_kinetics/test_util.py

@@ -1,8 +1,6 @@
-from NanoParticleTools.differential_kinetics import (get_parser,
-                                                     run_one_rate_eq,
-                                                     run_and_save_one,
-                                                     save_data_to_hdf5,
-                                                     load_data_from_hdf5)
+from NanoParticleTools.differential_kinetics.util import (
+    get_diff_kinetics_parser, run_one_rate_eq, run_and_save_one,
+    save_data_to_hdf5, load_data_from_hdf5, get_templates)
 from NanoParticleTools.species_data import Dopant
 
 import numpy as np
@@ -17,7 +15,8 @@ def hdf5_file(tmp_path):
 
 
 def test_parser():
-    args = get_parser().parse_args(['-n', '10', '-m', '2', '-o', 'test.h5'])
+    args = get_diff_kinetics_parser().parse_args(
+        ['-n', '10', '-m', '2', '-o', 'test.h5'])
     assert args.num_samples == 10
     assert args.max_dopants == 2
     assert args.output_file == 'test.h5'
@@ -27,6 +26,13 @@ def test_parser():
     assert args.dopants == ["Yb", "Er", "Tm", "Nd", "Ho", "Eu", "Sm", "Dy"]
 
 
+def test_get_templates():
+    args = get_diff_kinetics_parser().parse_args(
+        ['-n', '10', '-m', '2', '-o', 'test.h5'])
+    templates = get_templates(args)
+    assert len(list(templates)) == 10
+
+
 def test_run_one_rate_eq():
     dopants = [Dopant('Yb', 0.1), Dopant('Er', 0.1)]
     out_dict = run_one_rate_eq(dopants,
@@ -44,9 +50,10 @@ def test_run_and_save_one(hdf5_file):
     dopant_concs = [0.5, 0.2]
     run_and_save_one(dopants,
                      dopant_concs,
-                     0,
-                     3,
-                     hdf5_file,
+                     group_id=0,
+                     data_id=3,
+                     sample_id=3,
+                     file=hdf5_file,
                      excitation_wavelength=900,
                      excitation_power=1e5,
                      include_spectra=True)