AdeptML 1.0__tar.gz

adeptml-1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 ADAMS Lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

adeptml-1.0/PKG-INFO

@@ -0,0 +1,49 @@
+Metadata-Version: 2.1
+Name: AdeptML
+Version: 1.0
+Summary: A High-Level PyTorch based Library for Hybrid Physics-Informed Machine Learning Models
+Author: Manaswin Oddiraju
+Requires-Python: >=3.9.0,<3.13
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: docs
+Requires-Dist: Sphinx (==7.3.7) ; extra == "docs"
+Requires-Dist: jax (>=0.4.18,<0.5.0)
+Requires-Dist: joblib (>=1.3.2,<2.0.0)
+Requires-Dist: sphinx-rtd-theme (==2.0.0) ; extra == "docs"
+Requires-Dist: sphinxcontrib-napoleon (==0.7) ; extra == "docs"
+Requires-Dist: tensorboard (>=2.16.2,<3.0.0)
+Requires-Dist: torch (>=2.1.0,<3.0.0)
+Description-Content-Type: text/markdown
+
+# Auto-differentiable embedding of Physics and Torch Machine Learning (AdePT-ML):
+
+This is a convienience library built on top of PyTorch to enable easy integration and training of hybrid models involving physics and deep learning modules. 
+
+### Features
+1. Pre-defined Modules and configs for physics and MLP architectures.
+2. Physics Module accepts physics functions which return numpy arrays.
+3. Ability to input custom PyTorch nn classes.
+4. Ensemble Module allows easy integration of constituent physics and ML modules for simple inference and training.
+5. Auto-log training data with tensorboard.
+
+## Installation
+Installing with pip
+```
+pip install adeptml 
+```
+## Requirements (Automatically installed with pip): 
+1. PyTorch (https://pytorch.org/)
+2. Joblib (https://joblib.readthedocs.io/en/latest/) (For loading and saving model parameters)
+3. Tensorboard
+
+## Documentation:
+Visit our [Read the docs page](https://adept-ml.readthedocs.io/en/latest/)
+
+## Examples:
+Refer to the tests file. Additional examples will be added soon.
+
+

adeptml-1.0/README.md

@@ -0,0 +1,27 @@
+# Auto-differentiable embedding of Physics and Torch Machine Learning (AdePT-ML):
+
+This is a convienience library built on top of PyTorch to enable easy integration and training of hybrid models involving physics and deep learning modules. 
+
+### Features
+1. Pre-defined Modules and configs for physics and MLP architectures.
+2. Physics Module accepts physics functions which return numpy arrays.
+3. Ability to input custom PyTorch nn classes.
+4. Ensemble Module allows easy integration of constituent physics and ML modules for simple inference and training.
+5. Auto-log training data with tensorboard.
+
+## Installation
+Installing with pip
+```
+pip install adeptml 
+```
+## Requirements (Automatically installed with pip): 
+1. PyTorch (https://pytorch.org/)
+2. Joblib (https://joblib.readthedocs.io/en/latest/) (For loading and saving model parameters)
+3. Tensorboard
+
+## Documentation:
+Visit our [Read the docs page](https://adept-ml.readthedocs.io/en/latest/)
+
+## Examples:
+Refer to the tests file. Additional examples will be added soon.
+

adeptml-1.0/adeptml/__init__.py

@@ -0,0 +1,2 @@
+from adeptml.ensemble import HybridModel
+from adeptml.train_utils import train

adeptml-1.0/adeptml/configs.py

@@ -0,0 +1,87 @@
+import dataclasses
+from typing import Callable, Optional, Union
+import torch
+
+
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+else:
+    DEVICE = "cpu"
+
+
+@dataclasses.dataclass
+class MLPConfig:
+    """
+    Configuration class for the Multilayer Perceptron (MLP) model.
+
+    Attributes
+    ----------
+    layers : int
+        Total number of layers in the MLP (including hidden and output layers).
+
+    num_input_dim : int
+        Number of input dimensions to the MLP.
+
+    num_hidden_dim : int
+        Number of hidden dimensions in each hidden layer.
+
+    num_output_dim : int
+        Number of output dimensions from the MLP.
+
+    num_hidden_layers : int
+        Number of hidden layers in the MLP.
+
+    activation_functions : str
+        String representation of the activation functions used in the MLP.
+    """
+
+    num_input_dim: int
+    num_hidden_dim: int
+    num_output_dim: int
+    num_hidden_layers: int
+    activation_functions: str
+
+
+@dataclasses.dataclass
+class PhysicsConfig:
+    """
+    Configuration class for physics-related functions.
+
+    Attributes
+    ----------
+    forward_func : Callable[[torch.Tensor], torch.Tensor]
+        Forward function of the physics model.
+
+    jacobian_func : Callable[[torch.Tensor], torch.Tensor]
+        Function to compute the Jacobian of the physics model.
+    """
+
+    forward_func: Callable[[torch.Tensor], torch.Tensor]
+    jacobian_func: Callable[[torch.Tensor], torch.Tensor]
+
+
+ModelConfig = Union[MLPConfig, PhysicsConfig]
+
+
+@dataclasses.dataclass
+class HybridConfig:
+    """
+    Config for Ensemble Models.
+
+    Attributes
+    ----------
+    models: dict
+        Contains Modelname as keys and an instance of ModelConfigs as values.
+
+    model_inputs: dict
+        By default, the Ensemble model operates sequentially, using the output of the preceding model as input for the next.
+        Setting this dict to a non-empty value overrides that behavior.
+        Keys are model names; values are dicts specifying input customization.
+        Each inner dict holds model names as keys and specifies how to stack inputs:
+        - 'None' stacks the entire tensor.
+        - A list of ints stacks only specified dimensions.
+        Use "Input" if the input to this model matches the hybrid model's original input.
+    """
+
+    models: dict[str, ModelConfig | torch.nn.Module]
+    model_inputs: Optional[dict[str, dict[str, list[int] | None]]] = None

adeptml-1.0/adeptml/ensemble.py

@@ -0,0 +1,79 @@
+import torch
+import torch.nn
+from adeptml import configs, models
+
+
+class HybridModel(torch.nn.Module):
+    """
+    Torch Module for Serial Hybrid Physics Models.
+
+    Parameters
+    ----------
+    models_mlp : Torch module list of all MLP modules.
+    models_cnn : Torch module list of all CNN modules.
+    models_physics : Torch module list of all Physics modules.
+    unmodified_inputs : Indices of the inputs that are to be passed directly to the
+    model. These are appended to the outputs of the previous model.
+    architecture : Dict with key corresponding to model name and value being a model
+    config.
+    """
+
+    def __init__(self, config: configs.HybridConfig):
+        super(HybridModel, self).__init__()
+        self.models_nn = torch.nn.ModuleDict()
+        self.models_physics = {}
+        self.config = config
+        for model_name in config.models:
+            if isinstance(config.models[model_name], torch.nn.Module):
+                self.models_nn[model_name] = config.models[model_name]().to(
+                    configs.DEVICE
+                )
+            if isinstance(config.models[model_name], configs.PhysicsConfig):
+                self.models_physics[model_name] = models.Physics.apply
+            elif isinstance(config.models[model_name], configs.MLPConfig):
+                self.models_nn[model_name] = models.MLP(config.models[model_name]).to(
+                    configs.DEVICE
+                )
+            self.model_inputs = {}
+            self.interim_data = {}
+            if config.model_inputs:
+                to_save = []
+                for _, vals in config.model_inputs.items():
+                    to_save += list(vals.keys())
+                self.to_save = list(set(to_save))
+
+    def forward(self, x, phy_args=None):
+        """Function to run inference on the hybrid model."""
+        self.interim_data["Input"] = x
+        current_input = x
+        for model_name in self.config.models:
+            if self.config.model_inputs:
+                if model_name in self.config.model_inputs:
+                    input_tensors = []
+                    for input_model, dims in self.config.model_inputs[
+                        model_name
+                    ].items():
+                        if dims:
+                            input_tensors.append(
+                                self.interim_data[input_model][:, dims]
+                            )
+                        else:
+                            input_tensors.append(self.interim_data[input_model])
+                    current_input = torch.hstack(input_tensors)
+            if model_name in self.models_nn.keys():
+                cur_model = self.models_nn[model_name]
+                out = cur_model(current_input)
+            elif model_name in self.models_physics.keys():
+                cur_model = self.models_physics[model_name]
+                out = cur_model(
+                    current_input,
+                    self.config.models[model_name].forward_func,
+                    self.config.models[model_name].jacobian_func,
+                    phy_args,
+                )
+            if self.config.model_inputs:
+                if model_name in self.to_save:
+                    self.interim_data[model_name] = out
+            current_input = out
+
+        return out

adeptml-1.0/adeptml/models.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn
+from adeptml import configs
+
+ACTIVATIONS = {
+    "leakyrelu": torch.nn.LeakyReLU(),
+    "sigmoid": torch.nn.Sigmoid(),
+    "Tanh": torch.nn.Tanh(),
+}
+
+
+class MLP(torch.nn.Module):
+    """
+    Multilayer Perceptron (MLP) neural network model.
+
+    Attributes
+    ----------
+    config : Instance of MLPConfig dataclass.
+
+    Note
+    ----
+    This class implements a Multilayer Perceptron (MLP) neural network model.
+    It takes a configuration dictionary with parameters such as hidden layer size,
+    input and output dimensions, and the number of hidden layers.
+    """
+
+    def __init__(self, config: configs.MLPConfig):
+        super(MLP, self).__init__()
+        self.layers = torch.nn.ModuleList()
+        self.linear_in = torch.nn.Linear(config.num_input_dim, config.num_hidden_dim)
+        for _ in range(config.num_hidden_layers):
+            self.layers.append(
+                torch.nn.Linear(config.num_hidden_dim, config.num_hidden_dim)
+            )
+        self.linear_out = torch.nn.Linear(config.num_hidden_dim, config.num_output_dim)
+        self.nl1 = ACTIVATIONS[config.activation_functions]
+
+    def forward(self, x):
+        """
+        Forward pass of the MLP model.
+
+        :param torch.Tensor x: Input tensor.
+
+        :return: Output tensor.
+        :rtype: torch.Tensor
+        """
+        out = self.linear_in(x)
+        for i in range(len(self.layers) - 1):
+            net = self.layers[i]
+            out = self.nl1(net(out))
+        out = self.linear_out(out)
+        return out
+
+
+class Physics(torch.autograd.Function):
+    """Custom Autograd function to enable backpropagation on Custom Physics Models.
+
+    Attributes:
+    config: Instance of PhysicsConfig.
+    """
+
+    @staticmethod
+    def forward(ctx, x, forward_fun, jacobian_fun, args=None):
+        if args:
+            ctx.save_for_backward(x, *args)
+        else:
+            ctx.save_for_backward(x)
+        ctx.jacobian_fun = jacobian_fun
+        x = x.detach().cpu().numpy()
+        if args != None:
+            args = [tmp_args.detach().cpu().numpy() for tmp_args in args]
+            out = forward_fun(x, *args)
+            out = torch.Tensor(out).to(configs.DEVICE)
+            return out
+        else:
+            out = forward_fun(x)
+            out = torch.Tensor(out).to(configs.DEVICE)
+            return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input = ctx.saved_tensors[0]
+        args = ctx.saved_tensors[1:]
+        jacobian_fun = ctx.jacobian_fun
+        jac_final = None
+        if ctx.needs_input_grad[0]:
+            input = input.detach().cpu().numpy()
+            if args != None:
+                args = [tmp_args.detach().cpu().numpy() for tmp_args in args]
+                jac_final = jacobian_fun(input, *args)
+                jac_final = torch.Tensor(jac_final).to(configs.DEVICE)
+                jac_final = jac_final.reshape(input.shape[0], -1, input.shape[1])
+                grad_output = grad_output.reshape(input.shape[0], -1, 1)
+                grad_final = torch.matmul(grad_output, jac_final)
+                return grad_final, None, None, None
+            else:
+                jac_final = jacobian_fun(input)
+                jac_final = torch.Tensor(jac_final).to(configs.DEVICE)
+                grad_final = torch.zeros(input.shape[0], input.shape[1]).to(
+                    configs.DEVICE
+                )
+                grad_output = grad_output.reshape(input.shape[0], -1)
+                for i in range(grad_final.shape[0]):
+                    grad_final[i, :] = torch.matmul(
+                        grad_output[i, :].reshape(1, -1),
+                        jac_final[i].reshape(-1, input.shape[1]),
+                    )
+                return grad_final, None, None

adeptml-1.0/adeptml/train_utils.py

@@ -0,0 +1,129 @@
+import os
+import torch
+from torch.utils.tensorboard import SummaryWriter
+import numpy as np
+from adeptml import configs
+from adeptml import HybridModel
+
+
+def train_step(
+    model: HybridModel, optimizer: torch.optim.Optimizer, loss_fn, scheduler=None
+):
+    # Builds function that performs a step in the train loop
+    def train_step(x, y, args, test=False):
+        if not test:
+            yhat = model.forward(x, args)
+            loss = loss_fn(yhat, y)  # torch.mean(torch.abs(yhat-y))
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if scheduler:
+                scheduler.step()
+        else:
+            a = model.eval()
+            with torch.no_grad():
+                yhat = a(x, args)
+                loss = loss_fn(yhat, y)
+        return loss.item()
+
+    # Returns the function that will be called inside the train loop
+    return train_step
+
+
+def train(
+    model,
+    train_loader,
+    test_loader,
+    optimizer,
+    loss_fn,
+    scheduler,
+    filename,
+    epochs,
+    print_training_loss=True,
+    save_frequency=50,
+):
+    """Training Function.
+
+    Parameters
+    ----------
+    train_loader : torch.Torch_Dataloader
+        Torch Dataloader with training samples.
+
+    test_loader : torch.Torch_Dataloader
+        Torch Dataloader with validation samples.
+
+    optimizer : torch.optim.Optimizer
+        Initialized Torch Optimizer.
+
+    loss_fn : callable
+        Loss function for training.
+
+    scheduler : torch.optim.lr_scheduler
+        Learning rate scheduler.
+
+    filename : str
+        File name for saving the trained model.
+
+    epochs : int
+        Number of training epochs.
+
+    print_training_loss: bool
+        Option to toggle printing epoch loss.
+
+    save_frequency: int
+        Number of epochs per which to save the model parameters to disk.
+
+
+    Returns
+    -------
+        Trained Hybrid Model.
+    """
+
+    data_dir = os.path.join(os.getcwd(), f"Training_History_{filename}")
+    if not os.path.exists(data_dir):
+        os.system("mkdir %s" % data_dir)
+    try:
+        runs = max(
+            [int(f.name.split("_")[-1]) for f in os.scandir(data_dir) if f.is_dir()]
+        )
+    except:
+        runs = 0
+    current_data_dir = f"{data_dir}/run_{runs+1}"
+    cur_settings = ""
+    # for i in asdict(model.config) :
+    with SummaryWriter(log_dir=current_data_dir) as writer:
+        train_step_obj = train_step(model, optimizer, loss_fn, scheduler)
+        for epoch in range(epochs):
+            train_batch_losses = []
+            for data in train_loader:
+                x_batch = data[0]
+                y_batch = data[1]
+                if len(data) > 2:
+                    args = data[2:]
+                else:
+                    args = None
+                loss = train_step_obj(x_batch, y_batch, args)
+                train_batch_losses.append(loss)
+            writer.add_scalar("Loss/train", np.mean(train_batch_losses), epoch)
+            test_batch_losses = []
+            for data in test_loader:
+                x_batch = data[0]
+                y_batch = data[1]
+                if len(data) > 2:
+                    args = data[2:]
+                else:
+                    args = None
+                loss = train_step_obj(x_batch, y_batch, args, test=True)
+                test_batch_losses.append(loss)
+            writer.add_scalar("Loss/test", np.mean(test_batch_losses), epoch)
+            if print_training_loss:
+                print(
+                    f"Train Loss {np.mean(train_batch_losses)} Test Loss {np.mean(test_batch_losses)}"
+                )
+            if epoch % save_frequency == 0:
+                torch.save(
+                    model.state_dict(), "%s/Model_%d.pt" % (current_data_dir, epoch)
+                )
+        torch.save(model.state_dict(), current_data_dir + "/Model_final.pt")
+
+    return model

adeptml-1.0/pyproject.toml

@@ -0,0 +1,24 @@
+[tool.poetry]
+name = "AdeptML"
+version = "1.0"
+description = "A High-Level PyTorch based Library for Hybrid Physics-Informed Machine Learning Models"
+authors = ["Manaswin Oddiraju"]
+readme = "README.md"
+packages = [{include="adeptml"}]
+
+[tool.poetry.dependencies]
+python =">=3.9.0,<3.13"
+torch = "^2.1.0"
+jax = "^0.4.18"
+joblib = "^1.3.2"
+tensorboard = "^2.16.2"
+Sphinx = { version = "7.3.7", optional = true }
+sphinx-rtd-theme = { version = "2.0.0", optional = true }
+sphinxcontrib-napoleon = { version = "0.7", optional = true }
+
+[tool.poetry.extras]
+docs = ["Sphinx", "sphinx-rtd-theme", "sphinxcontrib-napoleon"]
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"