feat: parallelize viscosity calculation

Run 3 different temperatures in parallel.

Author: ltalirz

amorphouspy/src/amorphouspy/pipelines/viscosity.py

@@ -1,8 +1,8 @@
 """Multi-temperature viscosity pipeline.
 
 Sequentially cools a quenched glass to each target temperature via
-melt-quench and runs a Green-Kubo viscosity production simulation
-at each one.
+melt-quench, then runs Green-Kubo viscosity production simulations
+in parallel using an executor.
 """
 
 from __future__ import annotations
@@ -14,14 +14,86 @@
 from amorphouspy.fabrication.meltquench import melt_quench_simulation
 
 if TYPE_CHECKING:
+    from concurrent.futures import Executor
+
     import pandas as pd
     from ase import Atoms
 from amorphouspy.properties.viscosity import get_viscosity, viscosity_simulation
 
 logger = logging.getLogger(__name__)
 
 
+# ---------------------------------------------------------------------------
+# Helpers submitted to the executor
+# ---------------------------------------------------------------------------
+
+
+def _cool_and_run_viscosity(
+    structure: Atoms,
+    potential: pd.DataFrame,
+    temp_high: float,
+    temp_low: float,
+    heating_rate: float,
+    cooling_rate: float,
+    timestep: float,
+    n_timesteps: int,
+    n_print: int,
+    max_lag: int | None,
+    server_kwargs: dict[str, Any],
+) -> dict[str, Any]:
+    """Cool to *temp_low* via melt-quench, then run a viscosity simulation.
+
+    This function is submitted to the executor as a single unit of work
+    so that the melt-quench and the production run share the same worker.
+    """
+    mq_result = melt_quench_simulation(
+        structure=structure,
+        potential=potential,
+        temperature_high=float(temp_high),
+        temperature_low=float(temp_low),
+        timestep=1.0,
+        heating_rate=float(heating_rate),
+        cooling_rate=float(cooling_rate),
+        n_print=1000,
+        langevin=False,
+        server_kwargs=server_kwargs,
+    )
+    cooled_structure = mq_result["structure"]
+
+    visc_result = viscosity_simulation(
+        structure=cooled_structure,
+        potential=potential,
+        temperature_sim=float(temp_low),
+        timestep=float(timestep),
+        initial_production_steps=int(n_timesteps),
+        n_print=int(n_print),
+        langevin=False,
+        seed=12345,
+        server_kwargs=server_kwargs,
+    )
+
+    visc_data = get_viscosity(visc_result, timestep=float(timestep), max_lag=max_lag)
+
+    raw_lag = cast("list[float]", visc_data.get("lag_time_ps", []))
+    raw_visc = cast("list[float]", visc_data.get("viscosity_integral", []))
+
+    return {
+        "temperature": float(temp_low),
+        "viscosity": float(visc_data["viscosity"]),
+        "max_lag": float(visc_data["max_lag"]),
+        "simulation_steps": int(n_timesteps),
+        "lag_times_ps": downsample_log(raw_lag),
+        "viscosity_integral": downsample_log(raw_visc),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
 def run_viscosity_workflow(
+    executor: Executor,
     structure: Atoms,
     potential: pd.DataFrame,
     temperatures: list[float],
@@ -32,15 +104,19 @@ def run_viscosity_workflow(
     n_print: int = 1,
     max_lag: int | None = 1_000_000,
     server_kwargs: dict[str, Any] | None = None,
+    resource_dict: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
-    """Run viscosity analysis at multiple temperatures.
+    """Run viscosity analysis at multiple temperatures using an executor.
 
-    Starting from the quenched structure, the workflow sequentially cools
-    from the highest to the lowest requested temperature.  At each step a
-    melt-quench brings the structure to the target temperature, followed by
-    a Green-Kubo viscosity production run.
+    Each temperature is submitted as an independent task to *executor*.
+    Every task first cools the starting structure to the target temperature
+    via a melt-quench, then performs a Green-Kubo viscosity production run.
+    Because each task starts from the same *structure* and cools
+    independently, the tasks can run in parallel.
 
     Args:
+        executor: A ``concurrent.futures.Executor`` or executorlib executor
+            used to submit per-temperature tasks.
         structure: Quenched ASE Atoms from the melt-quench pipeline.
         potential: LAMMPS potential object.
         temperatures: Target temperatures (K) for viscosity runs.
@@ -51,67 +127,57 @@ def run_viscosity_workflow(
         n_print: Thermodynamic output frequency.
         max_lag: Maximum correlation lag (steps) for Green-Kubo.
         server_kwargs: LAMMPS server/resource configuration.
+        resource_dict: Optional resource dict forwarded to executorlib
+            executors (e.g. ``{"cores": 8}``).  Ignored for stdlib executors.
 
     Returns:
         Result dict with ``temperatures``, ``viscosities``, ``max_lag``,
         ``simulation_steps``, ``lag_times_ps``, and ``viscosity_integral``.
     """
     if server_kwargs is None:
         server_kwargs = {}
+    if resource_dict is None:
+        resource_dict = {}
 
     sorted_temps = sorted(temperatures, reverse=True)
-    logger.info("Viscosity temperatures (high→low): %s", sorted_temps)
-
+    logger.info("Submitting viscosity tasks for temperatures: %s", sorted_temps)
+
+    # Submit all temperatures in parallel — each cools independently
+    # from the starting structure down to its target temperature.
+    futures = {}
+    for temp in sorted_temps:
+        submit_kwargs: dict[str, Any] = {
+            "structure": structure,
+            "potential": potential,
+            "temp_high": 5000.0,
+            "temp_low": temp,
+            "heating_rate": heating_rate,
+            "cooling_rate": cooling_rate,
+            "timestep": timestep,
+            "n_timesteps": n_timesteps,
+            "n_print": n_print,
+            "max_lag": max_lag,
+            "server_kwargs": server_kwargs,
+        }
+        if resource_dict:
+            submit_kwargs["resource_dict"] = resource_dict
+        futures[temp] = executor.submit(_cool_and_run_viscosity, **submit_kwargs)
+
+    # Collect results in temperature order
     viscosities: list[float] = []
     all_max_lags: list[float] = []
     sim_steps: list[int] = []
     lag_times_ps: list[list[float]] = []
     viscosity_integral: list[list[float]] = []
 
-    structure_current = structure
-
-    for idx, temp in enumerate(sorted_temps):
-        temp_high = 5000.0 if idx == 0 else sorted_temps[idx - 1]
-
-        logger.info("Cooling from %.1f K to %.1f K", temp_high, temp)
-        mq_result = melt_quench_simulation(
-            structure=structure_current,
-            potential=potential,
-            temperature_high=float(temp_high),
-            temperature_low=float(temp),
-            timestep=1.0,
-            heating_rate=float(heating_rate),
-            cooling_rate=float(cooling_rate),
-            n_print=1000,
-            langevin=False,
-            server_kwargs=server_kwargs,
-        )
-        structure_current = mq_result["structure"]
-        logger.info("Cooled to %.1f K, %d atoms", temp, len(structure_current))
-
-        logger.info("Running viscosity simulation at %.1f K", temp)
-        visc_result = viscosity_simulation(
-            structure=structure_current,
-            potential=potential,
-            temperature_sim=float(temp),
-            timestep=float(timestep),
-            initial_production_steps=int(n_timesteps),
-            n_print=int(n_print),
-            langevin=False,
-            seed=12345,
-            server_kwargs=server_kwargs,
-        )
-
-        visc_data = get_viscosity(visc_result, timestep=float(timestep), max_lag=max_lag)
-        logger.info("Viscosity at %.1f K: %.3e Pa·s", temp, visc_data["viscosity"])
-
-        viscosities.append(float(visc_data["viscosity"]))
-        all_max_lags.append(float(visc_data["max_lag"]))
-        sim_steps.append(int(n_timesteps))
-        raw_lag_data = cast("list[float]", visc_data.get("lag_time_ps", []))
-        raw_visc_data = cast("list[float]", visc_data.get("viscosity_integral", []))
-        lag_times_ps.append(downsample_log(raw_lag_data))
-        viscosity_integral.append(downsample_log(raw_visc_data))
+    for temp in sorted_temps:
+        result = futures[temp].result()
+        logger.info("Viscosity at %.1f K: %.3e Pa·s", temp, result["viscosity"])
+        viscosities.append(result["viscosity"])
+        all_max_lags.append(result["max_lag"])
+        sim_steps.append(result["simulation_steps"])
+        lag_times_ps.append(result["lag_times_ps"])
+        viscosity_integral.append(result["viscosity_integral"])
 
     return {
         "temperatures": sorted_temps,

amorphouspy_api/src/amorphouspy_api/pipeline.py

@@ -88,22 +88,26 @@ def _run_structural_analysis(submission: JobSubmission, config: StructureAnalysi
 
 def _run_viscosity(submission: JobSubmission, config: ViscosityAnalysis, result: dict) -> dict:
     """Multi-temperature viscosity analysis on the quenched glass."""
+    from concurrent.futures import ProcessPoolExecutor
+
     from amorphouspy.pipelines.viscosity import run_viscosity_workflow
 
     from amorphouspy_api.executor import get_lammps_server_kwargs
 
-    return run_viscosity_workflow(
-        structure=result["melt_quench"]["final_structure"],
-        potential=result["structure_generation"]["potential"],
-        temperatures=config.temperatures,
-        heating_rate=int(submission.simulation.quench_rate * 100),
-        cooling_rate=int(submission.simulation.quench_rate),
-        timestep=config.timestep,
-        n_timesteps=config.n_timesteps,
-        n_print=config.n_print,
-        max_lag=config.max_lag,
-        server_kwargs=get_lammps_server_kwargs(),
-    )
+    with ProcessPoolExecutor() as exe:
+        return run_viscosity_workflow(
+            executor=exe,
+            structure=result["melt_quench"]["final_structure"],
+            potential=result["structure_generation"]["potential"],
+            temperatures=config.temperatures,
+            heating_rate=int(submission.simulation.quench_rate * 100),
+            cooling_rate=int(submission.simulation.quench_rate),
+            timestep=config.timestep,
+            n_timesteps=config.n_timesteps,
+            n_print=config.n_print,
+            max_lag=config.max_lag,
+            server_kwargs=get_lammps_server_kwargs(),
+        )
 
 
 def _run_cte(submission: JobSubmission, config: CTEFluctuations | CTETemperatureScan, result: dict) -> dict: