musica 0.12.2__cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

musica/test/unit/test_state.py

@@ -0,0 +1,325 @@
+"""Tests for the State class."""
+from __future__ import annotations
+import pytest  # type: ignore # pylint: disable=import-error
+import numpy as np  # type: ignore # pylint: disable=import-error
+from musica.types import State, MICM
+import musica.mechanism_configuration as mc
+from musica.mechanism_configuration import Mechanism
+
+
+def create_test_mechanism() -> Mechanism:
+    """Helper function to create a test mechanism."""
+    # Chemical species
+    A = mc.Species(
+        name="A",
+        diffusion_coefficient_m2_s=20.3,
+        molecular_weight_kg_mol=13.2,
+        other_properties={
+            "__absolute tolerance": "1.0e-30"})
+    B = mc.Species(name="B")
+    C = mc.Species(name="C")
+    M = mc.Species(name="M", is_third_body=True)
+
+    # Chemical phases
+    gas = mc.Phase(name="gas", species=[A, B, C, M])
+
+    # Reactions
+    my_arrhenius = mc.Arrhenius(
+        name="my arrhenius",
+        A=32.1, B=-2.3, C=102.3, D=63.4, E=-1.3,
+        gas_phase=gas,
+        reactants=[B],
+        products=[C],
+        other_properties={"__irrelevant": "2"},
+    )
+
+    my_other_arrhenius = mc.Arrhenius(
+        name="my other arrhenius",
+        A=29.3, B=-1.5, Ea=101.2, D=82.6, E=-0.98,
+        gas_phase=gas,
+        reactants=[A],
+        products=[(1.2, B)]
+    )
+
+    my_troe = mc.Troe(
+        name="my troe",
+        gas_phase=gas,
+        k0_A=1.2e-12,
+        k0_B=167,
+        k0_C=3,
+        kinf_A=136,
+        kinf_B=5,
+        kinf_C=24,
+        Fc=0.9,
+        N=0.8,
+        reactants=[B, A],
+        products=[C],
+        other_properties={"__irrelevant": "2"},
+    )
+
+    my_ternary = mc.TernaryChemicalActivation(
+        name="my ternary chemical activation",
+        gas_phase=gas,
+        k0_A=32.1,
+        k0_B=-2.3,
+        k0_C=102.3,
+        kinf_A=63.4,
+        kinf_B=-1.3,
+        kinf_C=908.5,
+        Fc=1.3,
+        N=32.1,
+        reactants=[B, A],
+        products=[C],
+        other_properties={"__irrelevant": "2"},
+    )
+
+    my_branched = mc.Branched(
+        name="my branched",
+        gas_phase=gas,
+        reactants=[A],
+        alkoxy_products=[B],
+        nitrate_products=[C],
+        X=1.2e-4,
+        Y=167,
+        a0=0.15,
+        n=9,
+        other_properties={"__irrelevant": "2"},
+    )
+
+    my_tunneling = mc.Tunneling(
+        name="my tunneling",
+        gas_phase=gas,
+        reactants=[B],
+        products=[C],
+        A=123.45,
+        B=1200.0,
+        C=1.0e8,
+        other_properties={"__irrelevant": "2"},
+    )
+
+    my_surface = mc.Surface(
+        name="my surface",
+        gas_phase=gas,
+        gas_phase_species=A,
+        reaction_probability=2.0e-2,
+        gas_phase_products=[B, C],
+        other_properties={"__irrelevant": "2"},
+    )
+
+    photo_b = mc.Photolysis(
+        name="photo B",
+        gas_phase=gas,
+        reactants=[B],
+        products=[C],
+        scaling_factor=12.3,
+        other_properties={"__irrelevant": "2"},
+    )
+
+    my_emission = mc.Emission(
+        name="my emission",
+        gas_phase=gas,
+        products=[B],
+        scaling_factor=12.3,
+        other_properties={"__irrelevant": "2"},
+    )
+
+    my_first_order_loss = mc.FirstOrderLoss(
+        name="my first order loss",
+        gas_phase=gas,
+        reactants=[C],
+        scaling_factor=12.3,
+        other_properties={"__irrelevant": "2"},
+    )
+
+    user_defined = mc.UserDefined(
+        name="my user defined",
+        gas_phase=gas,
+        reactants=[A, B],
+        products=[(1.3, C)],
+        scaling_factor=12.3,
+        other_properties={"__irrelevant": "2"}
+    )
+
+    # Mechanism
+    return mc.Mechanism(
+        name="Full Configuration",
+        species=[A, B, C, M],
+        phases=[gas],
+        reactions=[my_arrhenius, my_other_arrhenius, my_troe, my_ternary,
+                   my_branched, my_tunneling, my_surface, photo_b,
+                   my_emission, my_first_order_loss, user_defined],
+        version=mc.Version(1, 0, 0),
+    )
+
+
+def get_test_solver(mech: Mechanism) -> MICM:
+    """Helper function to create a test solver."""
+    return MICM(mechanism=mech)
+
+
+def test_state_initialization():
+    """Test State initialization with various grid cell configurations."""
+    # Use the test mechanism
+    mech = create_test_mechanism()
+
+    # Create MICM instance
+    solver = get_test_solver(mech)
+
+    # Test with valid input
+    state = solver.create_state(number_of_grid_cells=1)
+    assert isinstance(state, State)
+
+    # Test with multiple grid cells
+    state_multi = solver.create_state(number_of_grid_cells=3)
+    assert isinstance(state_multi, State)
+
+    # Test with invalid input
+    with pytest.raises(ValueError, match="number_of_grid_cells must be greater than 0"):
+        solver.create_state(number_of_grid_cells=0)
+
+
+def test_set_get_concentrations():
+    """Test setting and getting concentrations."""
+    # Use the test mechanism
+    mech = create_test_mechanism()
+    solver = get_test_solver(mech)
+
+    # Test single grid cell
+    state = solver.create_state(number_of_grid_cells=1)
+    concentrations = {"A": 1.0, "B": 2.0, "C": 3.0}
+    state.set_concentrations(concentrations)
+    result = state.get_concentrations()
+    assert result["A"][0] == 1.0
+    assert result["B"][0] == 2.0
+    assert result["C"][0] == 3.0
+
+    # Test multiple grid cells
+    state_multi = solver.create_state(number_of_grid_cells=2)
+    concentrations_multi = {"A": [1.0, 2.0], "B": [3.0, 4.0], "C": [5.0, 6.0]}
+    state_multi.set_concentrations(concentrations_multi)
+    result_multi = state_multi.get_concentrations()
+    assert result_multi["A"] == [1.0, 2.0]
+    assert result_multi["B"] == [3.0, 4.0]
+    assert result_multi["C"] == [5.0, 6.0]
+
+    # Test invalid species
+    with pytest.raises(ValueError, match="Species D not found in the mechanism"):
+        state.set_concentrations({"D": 1.0})
+
+    # Test invalid length
+    with pytest.raises(ValueError, match="must have length"):
+        state_multi.set_concentrations({"A": [1.0]})
+
+    # Test cannot set third-body species
+    with pytest.raises(ValueError, match="Species M not found in the mechanism"):
+        state.set_concentrations({"M": 1.0})
+
+
+def test_set_get_conditions():
+    """Test setting and getting environmental conditions."""
+    # Use the test mechanism
+    mech = create_test_mechanism()
+    solver = get_test_solver(mech)
+
+    # Test single grid cell
+    state = solver.create_state(number_of_grid_cells=1)
+    state.set_conditions(temperatures=300.0, pressures=101325.0)
+    conditions = state.get_conditions()
+    assert conditions["temperature"][0] == 300.0
+    assert conditions["pressure"][0] == 101325.0
+    assert np.isclose(conditions["air_density"][0], 40.9, rtol=0.1)  # Approximate value from ideal gas law
+
+    # Test multiple grid cells
+    state_multi = solver.create_state(number_of_grid_cells=2)
+    state_multi.set_conditions(
+        temperatures=[300.0, 310.0],
+        pressures=[101325.0, 101325.0],
+        air_densities=[40.9, 39.5]
+    )
+    conditions_multi = state_multi.get_conditions()
+    assert conditions_multi["temperature"] == [300.0, 310.0]
+    assert conditions_multi["pressure"] == [101325.0, 101325.0]
+    assert conditions_multi["air_density"] == [40.9, 39.5]
+
+    # Test invalid input length
+    with pytest.raises(ValueError, match="must be a list of length"):
+        state_multi.set_conditions(temperatures=[300.0])
+
+
+def test_set_get_user_defined_rate_parameters():
+    """Test setting and getting user-defined rate parameters."""
+    # Use the test mechanism which includes a user-defined reaction
+    mech = create_test_mechanism()
+    solver = get_test_solver(mech)
+
+    # Test single grid cell
+    state = solver.create_state(number_of_grid_cells=1)
+    params = {"EMIS.my emission": 1.0}
+    state.set_user_defined_rate_parameters(params)
+    result = state.get_user_defined_rate_parameters()
+    assert result["EMIS.my emission"][0] == 1.0
+
+    # Test multiple grid cells
+    state_multi = solver.create_state(number_of_grid_cells=2)
+    params_multi = {"PHOTO.photo B": [1.0, 2.0]}
+    state_multi.set_user_defined_rate_parameters(params_multi)
+    result_multi = state_multi.get_user_defined_rate_parameters()
+    assert result_multi["PHOTO.photo B"] == [1.0, 2.0]
+
+    # Test invalid parameter
+    with pytest.raises(ValueError, match="User-defined rate parameter invalid_param not found"):
+        state.set_user_defined_rate_parameters({"invalid_param": 1.0})
+
+    solver = get_test_solver(mech)
+
+    state = solver.create_state(number_of_grid_cells=1)
+
+    # Test species ordering
+    species_ordering = state.get_species_ordering()
+    assert isinstance(species_ordering, dict)
+    assert len(species_ordering) == 3  # A, B, C (M is third-body and not included)
+
+    # Dictionary style access
+    assert species_ordering["A"] >= 0
+    assert species_ordering["B"] >= 0
+    assert species_ordering["C"] >= 0
+
+    # Using get() method with default value
+    assert species_ordering.get("A", -1) >= 0  # returns value if key exists
+    assert species_ordering.get("Z", -1) == -1  # returns -1 since Z doesn't exist
+
+    # Test key membership
+    assert "A" in species_ordering
+    assert "B" in species_ordering
+    assert "C" in species_ordering
+    assert "M" not in species_ordering  # third-body not included
+
+    # Test parameter ordering
+    param_ordering = state.get_user_defined_rate_parameters_ordering()
+    assert isinstance(param_ordering, dict)
+    assert len(param_ordering) == 6
+
+    # Convert dict keys to list using list() function
+    param_names = list(param_ordering.keys())
+    assert len(param_names) == 6
+    assert isinstance(param_names[0], str)
+
+    # Alternative way using list comprehension
+    param_names_alt = [key for key in param_ordering]
+    assert param_names_alt == param_names
+
+    # Sort the keys if needed - useful for consistent ordering in tests
+    sorted_param_names = sorted(param_ordering.keys())
+    assert len(sorted_param_names) == 6
+    assert all(isinstance(name, str) for name in sorted_param_names)
+
+    # Verify all expected keys are present
+    expected_params = [
+        "PHOTO.photo B",
+        "EMIS.my emission",
+        "LOSS.my first order loss",
+        "SURF.my surface.effective radius [m]",
+        "SURF.my surface.particle number concentration [# m-3]",
+        "USER.my user defined"
+    ]
+    assert sorted(expected_params) == sorted(param_names)