Move augment and label into data package and remove data_operations

dfode_kit/cli/commands/augment.py

@@ -48,8 +48,7 @@ def add_command_parser(subparsers):
 def handle_command(args):
     import numpy as np
 
-    from dfode_kit.data_operations.augment_data import random_perturb
-    from dfode_kit.data.io_hdf5 import get_TPY_from_h5
+    from dfode_kit.data import get_TPY_from_h5, random_perturb
 
     print('Handling augment command')
     print(f'Loading data from h5 file: {args.h5_file}')

dfode_kit/cli/commands/label.py

@@ -34,7 +34,7 @@ def add_command_parser(subparsers):
 def handle_command(args):
     import numpy as np
 
-    from dfode_kit.data_operations import label_npy as label_main
+    from dfode_kit.data import label_npy as label_main
 
     try:
         labeled_data = label_main(

dfode_kit/data/__init__.py

@@ -16,6 +16,8 @@
     "nn_integrate",
     "integrate_h5",
     "calculate_error",
+    "random_perturb",
+    "label_npy",
 ]
 
 _ATTRIBUTE_MODULES = {
@@ -33,6 +35,8 @@
     "nn_integrate": ("dfode_kit.data.integration", "nn_integrate"),
     "integrate_h5": ("dfode_kit.data.integration", "integrate_h5"),
     "calculate_error": ("dfode_kit.data.integration", "calculate_error"),
+    "random_perturb": ("dfode_kit.data.augment", "random_perturb"),
+    "label_npy": ("dfode_kit.data.label", "label_npy"),
 }
 
 

dfode_kit/data/augment.py

@@ -0,0 +1,167 @@
+import time
+
+import cantera as ct
+import numpy as np
+
+
+def single_step(npstate, chem, time_step=1e-6):
+    gas = ct.Solution(chem)
+    T_old, P_old, Y_old = npstate[0], npstate[1], npstate[2:]
+    gas.TPY = T_old, P_old, Y_old
+    res_1st = [T_old, P_old] + list(gas.Y)
+    reactor = ct.IdealGasConstPressureReactor(gas, name='R1')
+    sim = ct.ReactorNet([reactor])
+
+    sim.advance(time_step)
+    new_TPY = [gas.T, gas.P] + list(gas.Y)
+    res_1st += new_TPY
+
+    return res_1st
+
+
+def random_perturb(
+    array: np.ndarray,
+    mech_path: str,
+    dataset: int,
+    heat_limit: bool,
+    element_limit: bool,
+    eq_ratio: float = 1,
+    frozenTem: float = 310,
+    alpha: float = 0.1,
+    gamma: float = 0.1,
+    cq: float = 10,
+    inert_idx: int = -1,
+    time_step: float = 1e-6,
+) -> np.ndarray:
+    array = array[array[:, 0] > frozenTem]
+
+    gas = ct.Solution(mech_path)
+    n_species = gas.n_species
+    maxT = np.max(array[:, 0])
+    minT = np.min(array[:, 0])
+    maxP = np.max(array[:, 1])
+    minP = np.min(array[:, 1])
+    maxN2 = np.max(array[:, -1])
+    minN2 = np.min(array[:, -1])
+
+    H_O_ratio_base = 2 * eq_ratio
+
+    num = 0
+    new_array = []
+    while num < dataset:
+        if heat_limit:
+            from dfode_kit.training.formation import formation_calculate
+
+            qdot_ = np.zeros_like(array[:, 0])
+            formation = formation_calculate(mech_path)
+            label_array = label(array, mech_path)
+            for i in range(label_array.shape[0]):
+                qdot_[i] = (
+                    -(
+                        formation
+                        * (
+                            label_array[i, 4 + n_species : 4 + 2 * n_species]
+                            - label_array[i, 2 : 2 + n_species]
+                        )
+                        / time_step
+                    ).sum()
+                )
+
+        for j in range(array.shape[0]):
+            test_tmp = np.copy(array[j])
+            k = 0
+            while True:
+                k += 1
+
+                test_r = np.copy(array[j])
+
+                test_tmp[0] = test_r[0] + (maxT - minT) * (2 * np.random.rand() - 1.0) * alpha
+                test_tmp[1] = test_r[1] + (maxP - minP) * (2 * np.random.rand() - 1.0) * alpha * 20
+                test_tmp[-1] = test_r[-1] + (maxN2 - minN2) * (2 * np.random.rand() - 1) * alpha
+                for i in range(2, array.shape[1] - 1):
+                    test_tmp[i] = np.abs(test_r[i]) ** (1 + (2 * np.random.rand() - 1) * alpha)
+                test_tmp[2:-1] = test_tmp[2:-1] / np.sum(test_tmp[2:-1]) * (1 - test_tmp[-1])
+
+                if heat_limit:
+                    label_test_tmp = np.array(single_step(test_tmp, mech_path))
+                    qdot_new_ = (
+                        -(
+                            formation
+                            * (
+                                label_test_tmp[4 + n_species : 4 + 2 * n_species]
+                                - label_test_tmp[2 : 2 + n_species]
+                            )
+                            / time_step
+                        ).sum()
+                    )
+
+                if element_limit:
+                    gas.TPY = test_tmp[0], test_tmp[1], test_tmp[2:]
+                    H_mole_fraction = gas.elemental_mole_fraction("H")
+                    O_mole_fraction = gas.elemental_mole_fraction("O")
+                    H_O_ratio = H_mole_fraction / O_mole_fraction
+
+                if heat_limit and element_limit:
+                    condition = (
+                        (minT * (1 - gamma)) <= test_tmp[0] <= (maxT * (1 + gamma))
+                        and (H_O_ratio_base * (1 - gamma)) <= H_O_ratio <= (H_O_ratio_base * (1 + gamma))
+                        and (qdot_new_ > 1 / cq * qdot_[j] and qdot_new_ < cq * qdot_[j])
+                    )
+                elif heat_limit and not element_limit:
+                    condition = (
+                        (minT * (1 - gamma)) <= test_tmp[0] <= (maxT * (1 + gamma))
+                        and (qdot_new_ > 1 / cq * qdot_[j] and qdot_new_ < cq * qdot_[j])
+                    )
+                elif not heat_limit and element_limit:
+                    condition = (
+                        (minT * (1 - gamma)) <= test_tmp[0] <= (maxT * (1 + gamma))
+                        and (H_O_ratio_base * (1 - gamma)) <= H_O_ratio <= (H_O_ratio_base * (1 + gamma))
+                    )
+                else:
+                    condition = (minT * (1 - gamma)) <= test_tmp[0] <= (maxT * (1 + gamma))
+
+                if condition or k > 20:
+                    break
+
+            if k <= 20:
+                new_array.append(test_tmp)
+
+        num = len(new_array)
+        print(num)
+
+    new_array = np.array(new_array)
+    new_array = new_array[np.random.choice(new_array.shape[0], size=dataset)]
+    unique_array = np.unique(new_array, axis=0)
+    print(unique_array.shape)
+    return unique_array
+
+
+def label(
+    array: np.ndarray,
+    mech_path: str,
+    time_step: float = 1e-06,
+) -> np.ndarray:
+    from dfode_kit.data.integration import advance_reactor
+
+    gas = ct.Solution(mech_path)
+    n_species = gas.n_species
+
+    labeled_data = np.empty((array.shape[0], 2 * n_species + 4))
+
+    reactor = ct.Reactor(gas, name='Reactor1', energy='off')
+    reactor_net = ct.ReactorNet([reactor])
+    reactor_net.rtol, reactor_net.atol = 1e-6, 1e-10
+
+    start_time = time.time()
+
+    for i, state in enumerate(array):
+        gas = advance_reactor(gas, state, reactor, reactor_net, time_step)
+        labeled_data[i, : 2 + n_species] = state[: 2 + n_species]
+        labeled_data[i, 2 + n_species :] = np.array([gas.T, gas.P] + list(gas.Y))
+
+    end_time = time.time()
+    total_time = end_time - start_time
+
+    print(f"Total time used: {total_time:.2f} seconds")
+
+    return labeled_data

dfode_kit/data/integration.py

@@ -1,5 +1,4 @@
 import h5py
-import torch
 import numpy as np
 import cantera as ct
 
@@ -28,8 +27,9 @@ def advance_reactor(gas, state, reactor, reactor_net, time_step):
     return gas
 
 
-@torch.no_grad()
 def load_model(model_path, device, model_class, model_layers):
+    import torch
+
     state_dict = torch.load(model_path, map_location='cpu')
 
     model = model_class(model_layers)
@@ -41,8 +41,9 @@ def load_model(model_path, device, model_class, model_layers):
     return model
 
 
-@torch.no_grad()
 def predict_Y(model, model_path, d_arr, mech, device):
+    import torch
+
     gas = ct.Solution(mech)
     n_species = gas.n_species
     expected_dims = 2 + n_species
@@ -79,7 +80,6 @@ def predict_Y(model, model_path, d_arr, mech, device):
     return next_Y
 
 
-@torch.no_grad()
 def nn_integrate(orig_arr, model_path, device, model_class, model_layers, time_step, mech, frozen_temperature=510):
     model = load_model(model_path, device, model_class, model_layers)
 

dfode_kit/data_operations/label_data.py → dfode_kit/data/label.py

@@ -1,45 +1,39 @@
 import time
-import numpy as np
+
 import cantera as ct
+import numpy as np
 
-from .h5_kit import advance_reactor
 
 def label_npy(
-    mech_path, 
+    mech_path,
     time_step,
     source_path,
 ):
-    # Load the chemical mechanism
+    from dfode_kit.data.integration import advance_reactor
+
     gas = ct.Solution(mech_path)
     n_species = gas.n_species
 
-    # Load the dataset containing initial states for the reactor
     test_data = np.load(source_path)
     print(f"Loaded dataset from: {source_path}")
     print(f"{test_data.shape=}")
 
-    # Prepare an array to store labeled data
     labeled_data = np.empty((test_data.shape[0], 2 * n_species + 4))
 
-    # Initialize Cantera reactor
     reactor = ct.Reactor(gas, name='Reactor1', energy='off')
     reactor_net = ct.ReactorNet([reactor])
     reactor_net.rtol, reactor_net.atol = 1e-6, 1e-10
 
-    # Start timing the simulation
     start_time = time.time()
 
-    # Process each state in the dataset
     for i, state in enumerate(test_data):
         gas = advance_reactor(gas, state, reactor, reactor_net, time_step)
-        labeled_data[i, :2 + n_species] = state[:2 + n_species]
-        labeled_data[i, 2 + n_species:] = np.array([gas.T, gas.P] + list(gas.Y))
+        labeled_data[i, : 2 + n_species] = state[: 2 + n_species]
+        labeled_data[i, 2 + n_species :] = np.array([gas.T, gas.P] + list(gas.Y))
 
-    # End timing of the simulation
     end_time = time.time()
     total_time = end_time - start_time
 
-    # Print the total time used and the path of the saved data
     print(f"Total time used: {total_time:.2f} seconds")
-    
-    return labeled_data
+
+    return labeled_data