Upload folder using huggingface_hub

README.md

@@ -0,0 +1,123 @@
+---
+license: unknown
+language:
+- en
+tags:
+- OneScience
+- fluid dynamics
+- topology optimization
+- Gaussian processes
+- physics-informed learning
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">GP_for_TO</span>
+  </strong>
+</p>
+
+# Model Overview
+
+GP_for_TO (Physics-informed GP-TO) is a physics-informed Gaussian-process framework for topology optimization developed by researchers at Northwestern University. It jointly optimizes material distributions and physical state variables in complex design domains without requiring an explicit mesh discretization.
+
+Paper: Simultaneous and Meshfree Topology Optimization with Physics-informed Gaussian Processes
+https://arxiv.org/abs/2408.03490
+
+# Model Description
+
+GP_for_TO uses a Gaussian-process architecture whose mean function is parameterized by a deep neural network. It provides a mesh-free approach to fluid topology optimization problems such as minimizing dissipated power in Stokes flow.
+
+# Usage
+
+## 1. OneCode
+
+Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:
+
+[Launch OneCode for one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Setup
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for import checks and small-scale pipeline validation, but full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version for the target cluster, is recommended.
+
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/GP_for_TO --local_dir ./GP_for_TO
+cd GP_for_TO
+```
+
+### Set Up the Runtime Environment
+
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Installation with uv is also supported
+pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+```bash
+# Activate Conda first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# Installation with uv is also supported
+pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+
+### Training Data
+The method does not rely on pregenerated topology structures or labeled physical fields. Instead, it samples spatial points within the design domain and constructs physics-informed training signals from the governing Stokes-flow equations, boundary conditions, dissipated-power objective, and material volume-fraction constraint. The paper evaluates four fluid topology optimization cases—Rugby, Pipe Bend, Diffuser, and Double Pipe—and compares the results against COMSOL solutions obtained with the SIMP method.
+
+### Training
+
+The default training configuration retains the original settings: `N_col_domain=10000`, `N_train_per_BC=25`, `num_iter=50000`, and `diff_method=Numerical`.
+
+```bash
+python scripts/train.py
+```
+### Model Weights
+This repository will provide GP_for_TO model weights in the `weights/` directory. The weights will be uploaded soon.
+
+## Inference
+
+```bash
+python scripts/inference.py
+```
+
+By default, inference loads `weight/gp_for_to.pt` and produces:
+
+```text
+result/inference/predictions.npz
+result/inference/inference_summary.json
+```
+
+`predictions.npz` contains five arrays: `x`, `u`, `v`, `p`, and `ro`.
+
+## Result Visualization
+
+```bash
+python scripts/result.py
+```
+
+
+# Official OneScience Resources
+
+| Platform | OneScience Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citations and License
+
+- Original GP_for_TO paper: [Simultaneous and Meshfree Topology Optimization with Physics-informed Gaussian Processes](https://arxiv.org/abs/2408.03490).
+- This repository retains the relevant source and attribution notices. Follow all applicable license requirements when using, modifying, or distributing its contents.

conf/config.yaml

@@ -0,0 +1,40 @@
+root:
+  problem: "doublepipe"
+  seed: 11
+  output_names: ["u", "v", "p", "ro"]
+
+  runtime:
+    onescience_src: null
+
+  data:
+    n_train_per_bc: 25
+    n_col_domain: 10000
+
+  model:
+    mean_function: "neural_network"
+    nn_layers_base: [64, 64, 64, 64, 64]
+    dtype: "float32"
+
+  training:
+    device: "auto"
+    gpu: 0
+    num_iter: 50000
+    lr_default: 0.001
+    diff_method: "Numerical"
+    checkpoint_path: "./weight/gp_for_to.pt"
+    output_dir: "./result/training"
+    plot_outputs: true
+    checkpoints: [1000, 10000, 20000, 30000, 40000, 50000]
+
+  inference:
+    device: "auto"
+    gpu: 0
+    checkpoint_path: "./weight/gp_for_to.pt"
+    output_dir: "./result/inference"
+    n_col_domain: 10000
+
+  fake_data:
+    problem: "doublepipe"
+    n_train_per_bc: 25
+    n_col_domain: 10000
+    output_dir: "./data/fake"

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"physics-informed-topology-optimization"}

model/__init__.py

@@ -0,0 +1,4 @@
+from .gpregression import GPR
+from .gpplus_model import GPPLUS
+from .horseshoe import LogHalfHorseshoePrior
+from .mollified_uniform import MollifiedUniformPrior

model/gpplus_model.py

@@ -0,0 +1,778 @@
+import numpy as np
+import torch
+import math
+import gpytorch
+from gpytorch.constraints import Positive
+from gpytorch.priors import NormalPrior
+from gpytorch.distributions import MultivariateNormal
+from onescience.utils.GP_TO.plot_latenth import plot_sep
+from .gpregression import GPR
+import gpytorch.kernels as kernels
+from .mollified_uniform import MollifiedUniformPrior
+from pandas import DataFrame
+import matplotlib.pyplot as plt
+from torch.nn.parameter import Parameter
+import torch.nn.functional as F 
+
+def setlevels(X, qual_index = None, return_label = False):
+    labels = []
+    if qual_index == []:
+        return X
+    if qual_index is None:
+        qual_index = list(range(X.shape[-1]))
+    # if type(X) == np.ndarray:
+    #     temp = torch.from_numpy(X).detach().clone()
+    temp = np.copy(X)
+    if type(X) == torch.Tensor:
+        temp = X.clone()
+    if temp.ndim > 1:
+        for j in qual_index:
+            l = np.sort(np.unique(temp[..., j])).tolist()
+            labels.append(l)
+            #l =  torch.unique(temp[..., j], sorted = True).tolist()
+            temp[..., j] = torch.tensor([*map(lambda m: l.index(m),temp[..., j])])
+    else:
+            l = torch.unique(temp, sorted = True)
+            temp = torch.tensor([*map(lambda m: l.tolist().index(m), temp)])
+    
+    
+    if temp.dtype == object:
+        temp = temp.astype(float)
+        if type(X) == np.ndarray:
+            temp = torch.from_numpy(temp)
+        
+        if return_label:
+            return temp, labels
+        else:
+            return temp
+    else:
+        if type(X) == np.ndarray:
+            temp = torch.from_numpy(temp)
+        if return_label:
+            return temp, labels
+        else:
+            return temp
+
+class GPPLUS(GPR):
+    """The latent Map GP regression model (LMGP) which extends GPs to handle categorical inputs.
+
+    :note: Binary categorical variables should not be treated as qualitative inputs. There is no 
+        benefit from applying a latent variable treatment for such variables. Instead, treat them
+        as numerical inputs.
+
+    :param train_x: The training inputs (size N x d). Qualitative inputs needed to be encoded as 
+        integers 0,...,L-1 where L is the number of levels. For best performance, scale the 
+        numerical variables to the unit hypercube.
+    """
+    def __init__(
+        self,
+        #transformation_of_A_parameters:str,
+        train_x:torch.Tensor,
+        train_y:torch.Tensor,
+        collocation_x:torch.Tensor,
+        qual_ind_lev = {},
+        multiple_noise = False,
+        lv_dim:int=2,
+        quant_correlation_class:str='Rough_RBF',
+        noise:float=5e-8,
+        fix_noise:bool=True,
+        lb_noise:float=1e-8,
+        NN_layers:list = [],
+        name_output:str='u',
+        encoding_type = 'one-hot',
+        uniform_encoding_columns = 2,
+        lv_columns = [] ,
+        basis='neural_network',
+        NN_layers_base=[4,4],
+        basis_function_size=None,  
+        device="cpu",
+        dtype= torch.float32
+    ) -> None:
+    
+        tkwargs = {}  # or dict()
+        tkwargs['dtype'] = dtype
+        tkwargs['device'] =device
+
+        qual_index = list(qual_ind_lev.keys())
+        all_index = set(range(train_x.shape[-1]))
+        quant_index = list(all_index.difference(qual_index))
+        num_levels_per_var = list(qual_ind_lev.values())
+        #------------------- lm columns --------------------------
+        lm_columns = list(set(qual_index).difference(lv_columns))
+        if len(lm_columns) > 0:
+            qual_kernel_columns = [*lv_columns, lm_columns]
+        else:
+            qual_kernel_columns = lv_columns
+
+        #########################
+        if len(qual_index) > 0:
+            train_x = setlevels(train_x, qual_index=qual_index)
+        #
+        if multiple_noise:
+            noise_indices = list(range(0,num_levels_per_var[0]))
+        else:
+            noise_indices = []
+
+        if len(qual_index) == 1 and num_levels_per_var[0] < 2:
+            temp = quant_index.copy()
+            temp.append(qual_index[0])
+            quant_index = temp.copy()
+            qual_index = []
+            lv_dim = 0
+        elif len(qual_index) == 0:
+            lv_dim = 0
+
+        quant_correlation_class_name = quant_correlation_class
+
+        if len(qual_index) == 0:
+            lv_dim = 0
+
+        if quant_correlation_class_name == 'Rough_RBF':
+            quant_correlation_class = 'RBFKernel'
+        
+        if quant_correlation_class_name == 'Matern32Kernel':
+            quant_correlation_class = 'Matern32Kernel'
+        
+        if quant_correlation_class_name == 'Matern52Kernel':
+            quant_correlation_class = 'Matern52Kernel'
+
+        if quant_correlation_class_name == 'Matern12Kernel':
+            quant_correlation_class = 'Matern12Kernel'
+
+        if len(qual_index) > 0:
+            ####################### Defined multiple kernels for seperate variables ###################
+            qual_kernels = []
+            for i in range(len(qual_kernel_columns)):
+                qual_kernels.append(kernels.RBFKernel(
+                    active_dims=torch.arange(lv_dim) + lv_dim * i) )
+                qual_kernels[i].initialize(**{'lengthscale':1.0})
+                qual_kernels[i].raw_lengthscale.requires_grad_(False)
+
+        if len(quant_index) == 0:
+            correlation_kernel = qual_kernels[0]
+            for i in range(1, len(qual_kernels)):
+                correlation_kernel *= qual_kernels[i]
+        else:
+            try:
+                quant_correlation_class = getattr(kernels,quant_correlation_class)
+            except:
+                raise RuntimeError(
+                    "%s not an allowed kernel" % quant_correlation_class
+                )
+    
+            if quant_correlation_class_name == 'RBFKernel':
+                quant_kernel = quant_correlation_class(
+                    ard_num_dims=len(quant_index),
+                    active_dims=len(qual_kernel_columns) * lv_dim+torch.arange(len(quant_index)),
+                    lengthscale_constraint= Positive(transform= torch.exp,inv_transform= torch.log)
+                )
+            elif quant_correlation_class_name == 'Rough_RBF':
+                quant_kernel = quant_correlation_class(
+                    ard_num_dims=len(quant_index),
+                    active_dims=len(qual_kernel_columns)*lv_dim+torch.arange(len(quant_index)),
+                    lengthscale_constraint= Positive(transform = lambda x: 2.0**(-0.5) * torch.pow(10,-x/2),inv_transform= lambda x: -2.0*torch.log10(x/2.0))
+                )
+
+
+            elif quant_correlation_class_name == 'Matern12Kernel':
+                quant_kernel = quant_correlation_class(
+                    ard_num_dims=len(quant_index),
+                    active_dims=len(qual_kernel_columns)*lv_dim+torch.arange(len(quant_index)),
+                    lengthscale_constraint= Positive(transform= lambda x: 2.0**(-0.5) * torch.pow(10,-x/2),inv_transform= lambda x: -2.0*torch.log10(x/2.0))
+                )
+            
+            elif quant_correlation_class_name == 'Matern32Kernel':
+                quant_kernel = quant_correlation_class(
+                    ard_num_dims=len(quant_index),
+                    active_dims=len(qual_kernel_columns)*lv_dim+torch.arange(len(quant_index)),
+                    #lengthscale_constraint= Positive(transform= torch.exp,inv_transform= torch.log)
+                    lengthscale_constraint= Positive(transform= lambda x: 2.0**(-0.5) * torch.pow(10,-x/2),inv_transform= lambda x: -2.0*torch.log10(x/2.0))             
+                )
+
+            elif quant_correlation_class_name == 'Matern52Kernel':
+                quant_kernel = quant_correlation_class(
+                    ard_num_dims=len(quant_index),
+                    active_dims=len(qual_kernel_columns)*lv_dim+torch.arange(len(quant_index)),
+                    #lengthscale_constraint= Positive(transform= torch.exp,inv_transform= torch.log)  
+                    lengthscale_constraint= Positive(transform= lambda x: 2.0**(-0.5) * torch.pow(10,-x/2),inv_transform= lambda x: -2.0*torch.log10(x/2.0))       
+                )
+
+            if quant_correlation_class_name == 'RBFKernel':
+                
+                quant_kernel.register_prior(
+                    'lengthscale_prior', MollifiedUniformPrior(math.log(0.1),math.log(10)),'raw_lengthscale'
+                )
+                
+            elif quant_correlation_class_name == 'Rough_RBF':
+                quant_kernel.register_prior(
+                    'lengthscale_prior',NormalPrior(-3.0,3.0),'raw_lengthscale'
+                )
+
+            elif quant_correlation_class_name == 'Matern12Kernel':
+                quant_kernel.register_prior(
+                    #'lengthscale_prior', MollifiedUniformPrior(math.log(0.1),math.log(10)),'raw_lengthscale'
+                    'lengthscale_prior',NormalPrior(-3.0,3.0),'raw_lengthscale'
+                )
+
+            elif quant_correlation_class_name == 'Matern32Kernel':
+                quant_kernel.register_prior(
+                    #'lengthscale_prior', MollifiedUniformPrior(math.log(0.1),math.log(10)),'raw_lengthscale'
+                    'lengthscale_prior',NormalPrior(-3.0,3.0),'raw_lengthscale'
+                )
+
+            elif quant_correlation_class_name == 'Matern52Kernel':
+                quant_kernel.register_prior(
+                    #'lengthscale_prior', MollifiedUniformPrior(math.log(0.1),math.log(10)),'raw_lengthscale'
+                    'lengthscale_prior',NormalPrior(-3.0,3.0),'raw_lengthscale'
+                )
+            
+            if len(qual_index) > 0:
+                temp = qual_kernels[0]
+                for i in range(1, len(qual_kernels)):
+                    temp *= qual_kernels[i]
+                correlation_kernel = temp*quant_kernel #+ qual_kernel + quant_kernel
+            else:
+                correlation_kernel = quant_kernel
+
+        super(GPPLUS,self).__init__(
+            train_x=train_x,train_y=train_y,noise_indices=noise_indices,
+            correlation_kernel=correlation_kernel,
+            noise=noise,fix_noise=fix_noise,lb_noise=lb_noise
+        )
+
+        # register index and transforms
+        self.register_buffer('quant_index',torch.tensor(quant_index))
+        self.register_buffer('qual_index',torch.tensor(qual_index))
+
+
+        self.num_levels_per_var = num_levels_per_var
+        self.lv_dim = lv_dim
+        self.uniform_encoding_columns = uniform_encoding_columns
+        self.encoding_type = encoding_type
+        self.perm =[]
+        self.zeta = []
+        self.perm_dict = []
+        self.A_matrix = []
+        self.collocation_x = collocation_x ####### ADDED
+        self.alpha = 1.0
+        self.beta = 20.0
+        self.covar_inv = None
+        self.omega = 3 #3.2
+        self.name_output = name_output
+        self.chol_decomp = None
+        self.g_uvp = None
+        self.k_xX = None
+
+        if len(qual_kernel_columns) > 0:
+            for i in range(len(qual_kernel_columns)):
+                if type(qual_kernel_columns[i]) == int:
+                    num = self.num_levels_per_var[qual_index.index(qual_kernel_columns[i])]
+                    cat = [num]
+                else:
+                    cat = [self.num_levels_per_var[qual_index.index(k)] for k in qual_kernel_columns[i]]
+                    num = sum(cat)
+
+                zeta, perm, perm_dict = self.zeta_matrix(num_levels=cat, lv_dim = self.lv_dim)
+                self.zeta.append(zeta.to(**tkwargs))
+                self.perm.append(perm.to(**tkwargs))
+                self.perm_dict.append(perm_dict)       
+
+                model_temp = FFNN(self, input_size= num, num_classes=lv_dim, 
+                    layers = NN_layers, name = str(qual_kernel_columns[i])).to(**tkwargs)
+                self.A_matrix.append(model_temp.to(**tkwargs))
+
+        self.basis=basis
+        i=0
+        if self.basis=='single':
+            self.mean_module = gpytorch.means.ConstantMean(prior=NormalPrior(0.,1.))
+            self.mean_module.constant.data = torch.tensor([0.0])  # Set the desired value
+            self.mean_module.constant.requires_grad = False  # Fix the hyperparameter
+        elif self.basis=='multiple_constant':
+            if basis_function_size is None:
+                basis_function_size=train_x.shape[1]-1
+            self.num_sources=int(torch.max(train_x[:,-1]))
+            for i in range(self.num_sources +1):
+                if i==0:
+                    setattr(self,'mean_module_'+str(i), gpytorch.means.ZeroMean())
+
+                else:
+                    #Constant 
+                    setattr(self,'mean_module_'+str(i), gpytorch.means.ConstantMean(prior=NormalPrior(0.,.3))) 
+                
+        elif self.basis=='multiple_polynomial':
+            if basis_function_size is None:
+                basis_function_size=train_x.shape[1]-1
+            self.num_sources=int(torch.max(train_x[:,-1]))
+            for i in range(self.num_sources +1):
+                if i==0:
+                    setattr(self,'mean_module_'+str(i), gpytorch.means.ZeroMean())
+                else:
+                    setattr(self,'mean_module_'+str(i), LinearMean_with_prior(input_size=basis_function_size, batch_shape=torch.Size([]), bias=True)) 
+        elif self.basis=='neural_network':
+            ############################################### One NN for ALL
+            if len(qual_index) == 0:
+                setattr(self,'mean_module_NN_All', FFNN_for_Mean(self, input_size= train_x.shape[1], num_classes=4,layers =NN_layers_base, name = str('mean_module_'+str(i)+'_'))) 
+            else:
+                setattr(self,'mean_module_NN_All', FFNN_for_Mean(self, input_size= train_x.shape[1]-len(qual_index)+2, num_classes=1, layers =NN_layers_base, name = str('mean_module_'+str(i)+'_'))) 
+        elif self.basis=='M3':
+            setattr(self,'mean_module_NN_All', NetworkM4(input_dim = train_x.shape[1], output_dim=3, layers = NN_layers_base)) 
+         
+
+        # Fix the hyperparameter value
+        self.covar_module.base_kernel.raw_lengthscale.data = torch.tensor([self.omega, self.omega], dtype=torch.float32)  # Set the desired value
+        self.covar_module.base_kernel.raw_lengthscale.requires_grad = False  # Fix the hyperparameter
+
+        self.covar_module.raw_outputscale.data = torch.tensor(0.541)  # Set the desired value
+        self.covar_module.raw_outputscale.requires_grad = False  # Fix the hyperparameter
+
+    def forward(self,x:torch.Tensor) -> MultivariateNormal:
+        x_forward_raw=x.clone()
+        nd_flag = 0
+        if x.dim() > 2:
+            xsize = x.shape
+            x = x.reshape(-1, x.shape[-1])
+            nd_flag = 1
+
+        if len(self.qual_kernel_columns) > 0:
+            embeddings = []
+            for i in range(len(self.qual_kernel_columns)):
+                temp= self.transform_categorical(x=x[:,self.qual_kernel_columns[i]].clone().type(torch.int64), 
+                    perm_dict = self.perm_dict[i], zeta = self.zeta[i])
+                embeddings.append(self.A_matrix[i](temp))
+
+            x= torch.cat([embeddings[0],x[...,self.quant_index]],dim=-1)
+
+        if nd_flag == 1:
+            x = x.reshape(*xsize[:-1], -1)
+
+#################### Multiple bases (General Case) ####################################  
+        def multi_mean(x,x_forward_raw):
+            mean_x=torch.zeros_like(x[:,-1])
+            if self.basis=='single':
+                mean_x=self.mean_module(x)
+            elif self.basis=='multiple_constant':
+                for i in range(len(mean_x)):
+                    qq=int(x_forward_raw[i,-1])
+                    mean_x[i]=getattr(self,'mean_module_'+str(qq))(torch.tensor(x[i,-1].clone()).reshape(-1,1))
+            elif self.basis=='multiple_polynomial':
+                for i in range(len(mean_x)):
+                    qq=int(x_forward_raw[i,-1])
+                    mean_x[i]=getattr(self,'mean_module_'+str(qq))(torch.cat((torch.tensor((x[i,-1].clone().double().reshape(-1,1))**2),torch.tensor(x[i,-1].clone().double()).reshape(-1,1)),1))
+
+            elif self.basis=='neural_network':
+                mean_x = getattr(self,'mean_module_NN_All')(x.clone())#.reshape(-1) #### FOR MULTIOUTPUT DELETE RESHAPE
+            
+            elif self.basis=='M3':
+                if hasattr(self, 'name_output'):
+                    mean_x = getattr(self,'mean_module_NN_All')(x.clone()) 
+                else:
+                    mean_x = getattr(self,'mean_module_NN_All')(x.clone()).reshape(-1)  
+
+            return mean_x 
+    ##########################################################################################
+        if self.name_output == 'u':
+            mean_x = multi_mean(x,x_forward_raw)[:,0].reshape(-1)
+        if self.name_output == 'v':
+            mean_x = multi_mean(x,x_forward_raw)[:,1].reshape(-1)
+        if self.name_output == 'p':
+            mean_x = multi_mean(x,x_forward_raw)[:,2].reshape(-1)
+        if self.name_output == 'ro':
+            mean_x = multi_mean(x,x_forward_raw)[:,3].reshape(-1)
+        covar_x = self.covar_module(x)
+
+        return MultivariateNormal(mean_x,covar_x)
+
+    def predict(self, Xtest,return_std=True, include_noise = True):
+        with torch.no_grad():
+            return super().predict(Xtest, return_std = return_std, include_noise= include_noise)
+            
+    def predict_with_grad(self, Xtest,return_std=True, include_noise = True):
+        return super().predict(Xtest, return_std = return_std, include_noise= include_noise)
+    
+    def noise_value(self):
+        noise = self.likelihood.noise_covar.noise.detach() * self.y_std**2
+        return noise
+
+    def visualize_latent(self, suptitle = None):
+        if len(self.qual_kernel_columns) > 0:
+            for i in range(len(self.qual_kernel_columns)):
+                zeta = self.zeta[i]
+                A = self.A_matrix[i]
+
+                positions = A(zeta)
+                level = torch.max(self.perm[i], axis = 0)[0].tolist()
+                perm = self.perm[i]
+                plot_sep(positions = positions, levels = level, perm = perm, constraints_flag=True, )
+
+        
+    def visualize_latent_position(self,lv_columns=None):                            
+        if len(self.qual_kernel_columns) > 0:
+            for i in range(len(self.qual_kernel_columns)):
+                zeta = self.zeta[i]
+                A = self.A_matrix[i]
+                positions = A(zeta)
+                if self.qual_kernel_columns[i]==lv_columns[0]:
+                    return positions       
+
+
+    def visualize_latent_position_simple(self, suptitle = None):
+        if len(self.qual_kernel_columns) > 0:
+            for i in range(len(self.qual_kernel_columns)):
+                zeta = self.zeta[i]
+                A = self.A_matrix[i]
+                positions = A(zeta)
+        return positions 
+
+
+    @classmethod
+    def show(cls):
+        plt.show()
+        
+    def get_params(self, name = None):
+        params = {}
+        print('###################Parameters###########################')
+        for n, value in self.named_parameters():
+             params[n] = value
+        if name is None:
+            print(params)
+            return params
+        else:
+            if name == 'Mean':
+                key = 'mean_module.constant'
+            elif name == 'Sigma':
+                key = 'covar_module.raw_outputscale'
+            elif name == 'Noise':
+                key = 'likelihood.noise_covar.raw_noise'
+            elif name == 'Omega':
+                for n in params.keys():
+                    if 'raw_lengthscale' in n and params[n].numel() > 1:
+                        key = n
+            print(params[key])
+            return params[key]
+
+    def get_latent_space(self):
+        if len(self.qual_index) > 0:
+            zeta = torch.tensor(self.zeta)
+            positions = self.nn_model(zeta)
+            return positions.detach()
+        else:
+            print('No categorical Variable, No latent positions')
+            return None
+        
+    def zeta_matrix(self,
+        num_levels:int,
+        lv_dim:int,
+        batch_shape=torch.Size()
+    ) -> None:
+
+        if any([i == 1 for i in num_levels]):
+            raise ValueError('Categorical variable has only one level!')
+
+        if lv_dim == 1:
+            raise RuntimeWarning('1D latent variables are difficult to optimize!')
+        
+        for level in num_levels:
+            if lv_dim > level - 0:
+                lv_dim = min(lv_dim, level-1)
+                raise RuntimeWarning(
+                    'The LV dimension can atmost be num_levels-1. '
+                    'Setting it to %s in place of %s' %(level-1,lv_dim)
+                )
+    
+        from itertools import product
+        levels = []
+        for l in num_levels:
+            levels.append(torch.arange(l))
+
+        perm = list(product(*levels))
+        perm = torch.tensor(perm, dtype=torch.int64)
+
+        #-------------Mapping-------------------------
+        perm_dic = {}
+        for i, row in enumerate(perm):
+            temp = str(row.tolist())
+            if temp not in perm_dic.keys():
+                perm_dic[temp] = i
+
+        #-------------One_hot_encoding------------------
+        for ii in range(perm.shape[-1]):
+            if perm[...,ii].min() != 0:
+                perm[...,ii] -= perm[...,ii].min()
+            
+        perm_one_hot = []
+        for i in range(perm.size()[1]):
+            perm_one_hot.append( torch.nn.functional.one_hot(perm[:,i]) )
+
+        perm_one_hot = torch.concat(perm_one_hot, axis=1)
+
+        return perm_one_hot, perm, perm_dic
+
+    
+    def transform_categorical(self, x:torch.Tensor,perm_dict = [], zeta = []) -> None:
+        if x.dim() == 1:
+            x = x.reshape(-1,1)
+        # categorical should start from 0
+        if self.training == False:
+            x = setlevels(x.cpu())
+        if self.encoding_type == 'one-hot':
+            index = [perm_dict[str(row.tolist())] for row in x]
+
+            if x.dim() == 1:
+                x = x.reshape(len(x),)
+
+            return zeta[index,:]  
+
+        elif self.encoding_type  == 'uniform':
+
+            temp2=np.random.uniform(0,1,(len(self.perm), self.uniform_encoding_columns))
+            dict={}
+            dict2={}
+
+            for i in range(0,self.perm.shape[0]):
+                dict[tuple((self.perm[i,:]).numpy())]=temp2[i,:]
+            
+            for i in range(0,x.shape[0]):
+                dict2[i]=dict[tuple((x[i]).numpy())]
+            
+            x_one_hot= torch.from_numpy(np.array(list(dict2.values())))
+        else:
+            raise ValueError ('Invalid type')
+
+                        
+        return x_one_hot
+
+
+##########################################################################################################################################################
+
+class LinearMean_with_prior(gpytorch.means.Mean):
+    def __init__(self, input_size, batch_shape=torch.Size(), bias=True):
+        super().__init__()
+        self.register_parameter(name="weights", parameter=torch.nn.Parameter(torch.randn(*batch_shape, input_size, 1)))
+        self.register_prior(name = 'weights_prior', prior=gpytorch.priors.NormalPrior(0.,1.), param_or_closure='weights')
+
+        if bias:
+            self.register_parameter(name="bias", parameter=torch.nn.Parameter(torch.randn(*batch_shape, 1)))
+            self.register_prior(name = 'bias_prior', prior=gpytorch.priors.NormalPrior(0.,1.), param_or_closure='bias')
+
+        else:
+            self.bias = None
+
+    def forward(self, x):
+        res = x.matmul(self.weights).squeeze(-1)
+        if self.bias is not None:
+            res = res + self.bias
+        return res
+
+
+###############################################################################################################################################################################################################################################################
+class FFNN(torch.nn.Module):
+    def __init__(self, lmgp, input_size, num_classes, layers,name):
+        super(FFNN, self).__init__()
+        # Our first linear layer take input_size, in this case 784 nodes to 50
+        # and our second linear layer takes 50 to the num_classes we have, in
+        # this case 10.
+        self.hidden_num = len(layers)
+        if self.hidden_num > 0:
+            self.fci = torch.nn.Linear(input_size, layers[0], bias=False) 
+            lmgp.register_parameter('fci', self.fci.weight)
+            lmgp.register_prior(name = 'latent_prior_fci', prior=gpytorch.priors.NormalPrior(0.,3.), param_or_closure='fci')
+
+            for i in range(1,self.hidden_num):
+                #self.h = nn.Linear(neuran[i-1], neuran[i])
+                setattr(self, 'h' + str(i), torch.nn.Linear(layers[i-1], layers[i], bias=False))
+                lmgp.register_parameter('h'+str(i), getattr(self, 'h' + str(i)).weight )
+                lmgp.register_prior(name = 'latent_prior'+str(i), prior=gpytorch.priors.NormalPrior(0.,3.), param_or_closure='h'+str(i))
+            
+            self.fce = torch.nn.Linear(layers[-1], num_classes, bias= False)
+            lmgp.register_parameter('fce', self.fce.weight)
+            lmgp.register_prior(name = 'latent_prior_fce', prior=gpytorch.priors.NormalPrior(0.,3.), param_or_closure='fce')
+        else:
+            self.fci = Linear_MAP(input_size, num_classes, bias = False)
+            lmgp.register_parameter(name, self.fci.weight)
+            lmgp.register_prior(name = 'latent_prior_'+name, prior=gpytorch.priors.NormalPrior(0,3) , param_or_closure=name)
+
+
+
+    def forward(self, x, transform = lambda x: x):
+        """
+        x here is the mnist images and we run it through fc1, fc2 that we created above.
+        we also add a ReLU activation function in between and for that (since it has no parameters)
+        I recommend using nn.functional (F)
+        """
+        if self.hidden_num > 0:
+            x = torch.tanh(self.fci(x))
+            for i in range(1,self.hidden_num):
+                x = torch.tanh( getattr(self, 'h' + str(i))(x) )
+            
+            x = self.fce(x)
+        else:
+            x = self.fci(x, transform)
+        return x
+    
+class FFNN_for_Mean(gpytorch.Module):
+    def __init__(self, lmgp, input_size, num_classes, layers, name):
+        super(FFNN_for_Mean, self).__init__()
+        self.dropout = torch.nn.Dropout(0.0)
+        # Our first linear layer take input_size, in this case 784 nodes to 50
+        # and our second linear layer takes 50 to the num_classes we have, in
+        # this case 10.
+        self.hidden_num = len(layers)
+        if self.hidden_num > 0:
+            self.fci = Linear_new(input_size, layers[0], bias=True, name='fci') 
+            for i in range(1,self.hidden_num):
+                setattr(self, 'h' + str(i), Linear_new(layers[i-1], layers[i], bias=True,name='h' + str(i)))
+            
+            self.fce = Linear_new(layers[-1], num_classes, bias=True,name='fce')
+        else:
+            self.fci = Linear_new(input_size, num_classes, bias=True,name='fci') #Linear_MAP(input_size, num_classes, bias = True)
+
+    def forward(self, x, transform = lambda x: x):
+        """
+        x here is the mnist images and we run it through fc1, fc2 that we created above.
+        we also add a ReLU activation function in between and for that (since it has no parameters)
+        I recommend using nn.functional (F)
+        """
+        if self.hidden_num > 0:
+            
+            x = torch.tanh(self.fci(x))
+            # x = self.dropout(x)
+            # x = self.fci(x)
+            for i in range(1,self.hidden_num):
+                x = torch.tanh( getattr(self, 'h' + str(i))(x) )
+
+            x = self.fce(x)
+            #x = torch.cat([x[...,:3] , torch.sigmoid(x[...,3]).unsqueeze(-1)] , dim = -1)
+        else:
+            x = self.fci(x)
+            #x = torch.cat([x[...,:3] , torch.tanh(x[...,3]).unsqueeze(-1)] , dim = -1)
+
+        return x
+    
+
+
+class NetworkM4(torch.nn.Module):
+    def __init__(self, input_dim = 2, output_dim = 1, layers = [40, 40, 40, 40], activation = 'tanh', collocation_x = []) -> None:
+        super(NetworkM4, self).__init__()
+        activation_list = {'tanh':torch.nn.Tanh(), 'Silu':torch.nn.SiLU(), 'Sigmoid':torch.nn.Sigmoid()}
+        activation = activation_list[activation]
+        self.dim = layers[0]
+  
+        self.U = torch.nn.Linear(input_dim, self.dim).to('cuda')
+        self.V = torch.nn.Linear(input_dim, self.dim).to('cuda')
+        self.H1 = torch.nn.Linear(input_dim, self.dim).to('cuda')
+        self.last= torch.nn.Linear(self.dim, output_dim).to('cuda')
+
+        self.collocation_x = collocation_x
+        self.alpha = 1.0
+        self.beta = 1.0
+        
+        l = torch.nn.ModuleList()
+        for _ in range(len(layers)):
+            l.append(torch.nn.Linear(self.dim, self.dim))
+            l.append(activation)
+        self.layers = torch.nn.Sequential(*l).to('cuda')
+
+    def forward(self, input):        
+        U = torch.nn.Tanh()(self.U(input))
+        V = torch.nn.Tanh()(self.V(input))
+        H = torch.nn.Tanh()(self.H1(input))
+
+        for layer in self.layers:
+            Z = layer(H)
+            H = (1-Z)*U + Z*V
+        
+        out = self.last(H)
+        return out
+
+class Linear_new(gpytorch.means.Mean):
+    r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`
+
+    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    Args:
+        in_features: size of each input sample
+        out_features: size of each output sample
+        bias: If set to ``False``, the layer will not learn an additive bias.
+            Default: ``True``
+
+    Shape:
+        - Input: :math:`(*, H_{in})` where :math:`*` means any number of
+          dimensions including none and :math:`H_{in} = \text{in\_features}`.
+        - Output: :math:`(*, H_{out})` where all but the last dimension
+          are the same shape as the input and :math:`H_{out} = \text{out\_features}`.
+
+    Attributes:
+        weight: the learnable weights of the module of shape
+            :math:`(\text{out\_features}, \text{in\_features})`. The values are
+            initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
+            :math:`k = \frac{1}{\text{in\_features}}`
+        bias:   the learnable bias of the module of shape :math:`(\text{out\_features})`.
+                If :attr:`bias` is ``True``, the values are initialized from
+                :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                :math:`k = \frac{1}{\text{in\_features}}`
+
+    Examples::
+
+        >>> m = nn.Linear(20, 30)
+        >>> input = torch.randn(128, 20)
+        >>> output = m(input)
+        >>> print(output.size())
+        torch.Size([128, 30])
+    """
+    __constants__ = ['in_features', 'out_features']
+    in_features: int
+    out_features: int
+    weight: torch.Tensor
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True, name=None,
+                 device=None, dtype=None) -> None:
+        # factory_kwargs = {'device': device, 'dtype': dtype}
+        # factory_kwargs=tkwargs
+        super(Linear_new, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+
+        self.name=str(name)
+        self.register_parameter(name=str(self.name)+'weight',  parameter= Parameter(torch.empty((out_features, in_features))))
+        self.register_prior(name =str(self.name)+ 'prior_m_weight_fci', prior=gpytorch.priors.NormalPrior(0.,1.), param_or_closure=str(self.name)+'weight')
+        
+        if bias:
+            self.register_parameter(name=str(self.name)+'bias',  parameter=Parameter(torch.empty(out_features)))
+            self.register_prior(name= str(self.name)+'prior_m_bias_fci', prior=gpytorch.priors.NormalPrior(0.,1.), param_or_closure=str(self.name)+'bias')
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:                                         
+        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
+        # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see
+        # https://github.com/pytorch/pytorch/issues/57109
+        torch.nn.init.kaiming_uniform_( getattr(self,str(self.name)+'weight'), a=math.sqrt(5))
+        if getattr(self,str(self.name)+'bias') is not None:
+            fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(getattr(self,str(self.name)+'weight'))
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            torch.nn.init.uniform_(getattr(self,str(self.name)+'bias'), -bound, bound)
+
+    def forward(self, input) -> torch.Tensor:
+        # return F.linear(input, self.weight, self.bias)
+        # print(getattr(self,str(self.name)+'weight'))
+
+        # return F.linear(input, getattr(self,str(self.name)+'weight').double(), getattr(self,str(self.name)+'bias').double())      ### Forced to Add .double() for NN in mean function
+        return F.linear(input, getattr(self,str(self.name)+'weight'), getattr(self,str(self.name)+'bias'))      ### Forced to Add .double() for NN in mean function
+
+    def extra_repr(self) -> str:
+        return 'in_features={}, out_features={}, bias={}'.format(
+            self.in_features, self.out_features, self.bias is not None
+        )
+
+class Linear_MAP(torch.nn.Linear):
+    def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None) -> None:
+        super().__init__(in_features, out_features, bias, device, dtype)
+        
+    def forward(self, input, transform = lambda x: x):
+        return F.linear(input,transform(self.weight), self.bias)
+
+
+

model/gpregression.py

@@ -0,0 +1,159 @@
+import torch
+import gpytorch
+import math
+from gpytorch.models import ExactGP
+from gpytorch import settings as gptsettings
+from gpytorch.constraints import GreaterThan,Positive
+from gpytorch.distributions import MultivariateNormal
+import gpytorch.kernels as kernels
+from .horseshoe import LogHalfHorseshoePrior
+from .mollified_uniform import MollifiedUniformPrior
+from gpytorch.priors import NormalPrior,LogNormalPrior
+from onescience.utils.GP_TO.transforms import softplus,inv_softplus
+
+from typing import List,Tuple,Union
+
+from typing import List, Union
+
+class GPR(ExactGP):
+    """Standard GP regression module for numerical inputs
+
+    :param train_x: The training inputs (size N x d). All input variables are expected
+        to be numerical. For best performance, scale the variables to the unit hypercube.
+    :type train_x: torch.Tensor
+    :param train_y: The training targets (size N)
+    :type train_y: torch.Tensor
+    :param correlation_kernel: Either a `gpytorch.kernels.Kernel` instance or one of the 
+        following strings - 'RBFKernel' (radial basis kernel), 'Matern52Kernel' (twice 
+        differentiable Matern kernel), 'Matern32Kernel' (first order differentiable Matern
+        kernel). If the former is specified, any hyperparameters to be estimated need to have 
+        associated priors for multi-start optimization. If the latter is specified, then 
+        the kernel uses a separate lengthscale for each input variable.
+    :type correlation_kernel: Union[gpytorch.kernels.Kernel,str]
+    :param noise: The (initial) noise variance.
+    :type noise: float, optional
+    :param fix_noise: Fixes the noise variance at the current level if `True` is specifed.
+        Defaults to `False`
+    :type fix_noise: bool, optional
+    :param lb_noise: Lower bound on the noise variance. Setting a higher value results in
+        more stable computations, when optimizing noise variance, but might reduce 
+        prediction quality. Defaults to 1e-6
+    :type lb_noise: float, optional
+    """
+    def __init__(
+        self,
+        train_x:torch.Tensor,
+        train_y:torch.Tensor,
+        correlation_kernel,
+        noise_indices:List[int],
+        noise:float=1e-4,
+        fix_noise:bool=False,
+        lb_noise:float=1e-12,
+    ) -> None:
+        # check inputs
+        if not torch.is_tensor(train_x):
+            raise RuntimeError("'train_x' must be a tensor")
+        if not torch.is_tensor(train_y):
+            raise RuntimeError("'train_y' must be a tensor")
+
+        if train_x.shape[0] != train_y.shape[0]:
+            raise RuntimeError("Inputs and output have different number of observations")
+        
+        # initializing likelihood
+        noise_constraint=GreaterThan(lb_noise,transform=torch.exp,inv_transform=torch.log)
+        
+        if len(noise_indices) == 0:
+            likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=noise_constraint)
+
+        y_mean = torch.tensor(0.0)
+        y_std = torch.tensor(1.0)
+        train_y_sc = (train_y-y_mean)/y_std
+
+        ExactGP.__init__(self, train_x,train_y_sc, likelihood)
+        
+        # registering mean and std of the raw response
+        self.register_buffer('y_mean',y_mean)
+        self.register_buffer('y_std',y_std)
+        self.register_buffer('y_scaled',train_y_sc)
+
+        self._num_outputs = 1
+
+        # initializing and fixing noise
+        if noise is not None:
+            self.likelihood.initialize(noise=noise)
+        
+        self.likelihood.register_prior('noise_prior',LogHalfHorseshoePrior(0.01,lb_noise),'raw_noise')
+        if fix_noise:
+            self.likelihood.raw_noise.requires_grad_(False)
+
+        if isinstance(correlation_kernel,str):
+            try:
+                correlation_kernel_class = getattr(kernels,correlation_kernel)
+                correlation_kernel = correlation_kernel_class(
+                    ard_num_dims = self.train_inputs[0].size(1),
+                    lengthscale_constraint=Positive(transform=torch.exp,inv_transform=torch.log),
+                )
+                correlation_kernel.register_prior(
+                    'lengthscale_prior',MollifiedUniformPrior(math.log(0.1),math.log(10)),'raw_lengthscale'
+                )
+            except:
+                raise RuntimeError(
+                    "%s not an allowed kernel" % correlation_kernel
+                )
+        elif not isinstance(correlation_kernel,gpytorch.kernels.Kernel):
+            raise RuntimeError(
+                "specified correlation kernel is not a `gpytorch.kernels.Kernel` instance"
+            )
+
+        self.covar_module = kernels.ScaleKernel(
+            base_kernel = correlation_kernel,
+            outputscale_constraint=Positive(transform=softplus,inv_transform=inv_softplus),
+        )
+        # register priors
+        self.covar_module.register_prior(
+            'outputscale_prior',LogNormalPrior(1e-6,1.),'outputscale'
+        )
+
+        
+
+    def forward(self,x:torch.Tensor)->MultivariateNormal:
+        mean_x = self.mean_module(x)
+        covar_x = self.covar_module(x)
+        return MultivariateNormal(mean_x,covar_x)
+    
+    def predict(
+        self,x:torch.Tensor,return_std:bool=False,include_noise:bool=False
+    )-> Union[torch.Tensor,Tuple[torch.Tensor]]:
+        """Returns the predictive mean, and optionally the standard deviation at the given points
+
+        :param x: The input variables at which the predictions are sought. 
+        :type x: torch.Tensor
+        :param return_std: Standard deviation is returned along the predictions  if `True`. 
+            Defaults to `False`.
+        :type return_std: bool, optional
+        :param include_noise: Noise variance is included in the standard deviation if `True`. 
+            Defaults to `False`.
+        :type include_noise: bool
+        """
+        self.eval()
+        with gptsettings.fast_computations(log_prob=False):
+            # determine if batched or not
+            ndim = self.train_targets.ndim
+            if ndim == 1:
+                output = self(x)
+            else:
+                # for batched GPs 
+                num_samples = self.train_targets.shape[0]
+                output = self(x.unsqueeze(0).repeat(num_samples,1,1))
+            
+            if return_std and include_noise:
+                output = self.likelihood(output)
+
+            out_mean = self.y_mean + self.y_std*output.mean
+
+            # standard deviation may not always be needed
+            if return_std:
+                out_std = output.variance.sqrt()*self.y_std
+                return out_mean,out_std
+
+            return out_mean  

model/horseshoe.py

@@ -0,0 +1,63 @@
+import math
+import torch
+from gpytorch.priors import Prior
+from torch.distributions import HalfCauchy,HalfNormal,constraints
+from torch.distributions.utils import broadcast_all
+from numbers import Number
+
+class LogHalfHorseshoePrior(Prior):
+    """Prior for the log-noise variance hyperparameter for GPs. 
+
+    This is parameterized by `scale` and `lb`. `lb` is the lower bound on the noise variance. 
+    The `scale` parameter is more important. The default value for `scale` - 0.01 - works well 
+    for deterministic and low-noise situations. A larger value may be need in noisy sitations 
+    and small training datasets. A larger scale implies more noisy data as the prior.
+
+    To change the scale for this prior for a model to say 0.1,
+        >>> model.likelihood.register(
+        >>>     'noise_prior',LogHalfHorseshoePrior(0.1,model.likelihood.noise_prior.lb),
+        >>>     'raw_noise'
+        >>> )
+
+    .. note::
+        The `log_prob` method is only approximate and unnormalized. There is no closed form
+        expression for the underlying horseshoe distribution. The lower and upper bounds on
+        its' density are, however, known. Here, we use the same approximate density value that 
+        the spearmint package uses.
+    
+    :param scale: scale parameter of the Horseshoe distribution
+    :type scale: float or torch.Tensor
+
+    :param lb: lower bound on the original scale. Defaults to 1e-6
+    :type lb: float or torch.Tensor, optional
+    """
+    arg_constraints = {"scale": constraints.positive,"lb":constraints.positive}
+    support = constraints.real
+    def __init__(self, scale, lb=1e-6,validate_args=None):
+        self.scale,self.lb = broadcast_all(scale,lb)
+        if isinstance(scale,Number):
+            batch_shape = torch.Size()
+        else:
+            batch_shape = self.scale.size()
+        super().__init__(batch_shape,validate_args=validate_args)
+    
+    def transform(self, x):
+        return self.lb + torch.exp(x)
+
+    def log_prob(self, X):
+        # first term is the density in the original scale
+        # the second term is for the transformation
+        return torch.log(torch.log(1+3*(self.scale / self.transform(X)) ** 2))+ X
+
+    def rsample(self, sample_shape=torch.Size([])):
+        local_shrinkage = HalfCauchy(1).rsample(self.scale.shape).to(self.lb)
+        param_sample = HalfNormal(local_shrinkage * self.scale).rsample(sample_shape).to(self.lb)
+        if len(self.lb) > 1:
+            param_sample[param_sample<self.lb[0]] = self.lb[0]
+        else:
+            param_sample[param_sample<self.lb] = self.lb
+        return param_sample.log()
+
+    def expand(self,expand_shape, _instance=None):
+        batch_shape = torch.Size(expand_shape)
+        return LogHalfHorseshoePrior(self.scale.expand(batch_shape))

model/mollified_uniform.py

@@ -0,0 +1,77 @@
+import math
+import torch
+from gpytorch.priors import Prior
+from torch.distributions import constraints,Uniform,Normal
+from torch.distributions.utils import broadcast_all
+from numbers import Number
+
+class MollifiedUniformPrior(Prior):
+    r"""Uniform distribution that is differentiable everywhere 
+
+    This is an approximation to the Uniform distribution which maintains differentiability by placing a 
+    Gaussian distribution over points away from the original support. The density for a single dimension is
+
+    .. math::
+
+        p(x) &= \frac{M}{\sqrt{2\pi\sigma^2}}\exp\left[-\frac{d(x)^2}{2\sigma^2}\right], \\
+        d(x) &= \begin{cases}
+            a-x & x < a \\
+            0 & a\leq x < b\\
+            x-b & x \geq b
+        \end{cases}
+
+    :param a: lower range (inclusive)
+    :type a: float or torch.Tensor
+
+    :param b: upper range (exclusive)
+    :type b: float or torch.Tensor
+
+    :param tail_sigma: Standard deviation of the Gaussian distributions on the tails. Lower values make the 
+        approximation closer to Uniform, but may increase the optimization effort. Defaults to 0.1
+    :type tail_sigma: float or torch.Tensor, optional
+
+    .. note::
+        The `rsample` method for this distribution returns uniformly distributed samples from the interval `[a,b)`, 
+        and **not** from the Mollified distribution. The `log_prob` method, however, returns the correct probability.
+        The purpose of the priors in this package is to generate initial starting points for optimizing
+        hyperparameters and for MAP estimation. 
+    """
+    arg_constraints = {'a':constraints.real,'b':constraints.real,'tail_sigma':constraints.positive}
+    support = constraints.real
+    def __init__(self,a,b,tail_sigma=0.1):
+        self.a,self.b,self.tail_sigma = broadcast_all(a,b,tail_sigma)
+        
+        if isinstance(a,Number) or isinstance(b,Number):
+            batch_shape = torch.Size()
+        else:
+            batch_shape = self.a.size()
+
+        super().__init__(batch_shape)
+    
+    @property
+    def mean(self):
+        return (self.a+self.b)/2
+    
+    @property
+    def _half_range(self):
+        return (self.b-self.a)/2
+
+    @property
+    def _log_normalization_constant(self):
+        return -torch.log(1+(self.b-self.a)/(math.sqrt(2*math.pi)*self.tail_sigma))
+
+    def log_prob(self,X):
+        # expression preserving gradients under automatic differentiation
+        tail_dist = ((X-self.mean).abs()-self._half_range).clamp(min=0)
+        return Normal(loc=torch.zeros_like(self.a),scale=self.tail_sigma).log_prob(tail_dist)+self._log_normalization_constant
+    
+    def rsample(self,sample_shape=torch.Size([])):
+        return Uniform(self.a,self.b).rsample(sample_shape).to(self.a)
+
+    def expand(self,expand_shape):
+        batch_shape = torch.Size(expand_shape)
+        return MollifiedUniformPrior(
+            self.a.expand(batch_shape),
+            self.b.expand(batch_shape),
+            self.tail_sigma.expand(batch_shape)
+        )

scripts/__init__.py

@@ -0,0 +1 @@
+

scripts/common.py

@@ -0,0 +1,155 @@
+import json
+import os
+import sys
+from pathlib import Path
+
+import torch
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+CONFIG_PATH = PROJECT_ROOT / "conf" / "config.yaml"
+PROBLEMS = ("doublepipe", "diffuser", "rugby", "pipebend")
+
+
+def ensure_onescience_path(explicit_src=None):
+    candidates = []
+    if explicit_src:
+        candidates.append(Path(explicit_src).expanduser())
+    if os.environ.get("ONESCIENCE_SRC"):
+        candidates.append(Path(os.environ["ONESCIENCE_SRC"]).expanduser())
+    for parent in (PROJECT_ROOT, *PROJECT_ROOT.parents):
+        candidates.append(parent / "refactor" / "onescience" / "src")
+        candidates.append(parent / "onescience" / "src")
+
+    for candidate in candidates:
+        if (candidate / "onescience").is_dir():
+            path = str(candidate.resolve())
+            if path not in sys.path:
+                sys.path.insert(0, path)
+            return candidate.resolve()
+    return None
+
+
+def load_config():
+    with CONFIG_PATH.open("r", encoding="utf-8") as f:
+        return yaml.safe_load(f)["root"]
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def select_device(section):
+    requested = section.get("device", "auto")
+    gpu = int(section.get("gpu", 0))
+    if requested == "auto":
+        return torch.device(f"cuda:{gpu}" if torch.cuda.is_available() else "cpu")
+    if requested.startswith("cuda") and not torch.cuda.is_available():
+        print(f"Requested {requested}, but CUDA is not available. Falling back to CPU.")
+        return torch.device("cpu")
+    return torch.device(requested)
+
+
+def dtype_from_config(name):
+    mapping = {
+        "float32": torch.float32,
+        "float": torch.float32,
+        "float64": torch.float64,
+        "double": torch.float64,
+    }
+    if name not in mapping:
+        raise ValueError(f"Unsupported dtype: {name}")
+    return mapping[name]
+
+
+def set_nested(cfg, section, key, value):
+    if value is not None:
+        cfg[section][key] = value
+
+
+def build_models(cfg, device, *, n_col_domain=None, n_train_per_bc=None, problem=None):
+    ensure_onescience_path(cfg.get("runtime", {}).get("onescience_src"))
+    from model import GPPLUS
+    from onescience.utils.GP_TO import get_data_fluid, set_seed
+
+    problem = problem or cfg["problem"]
+    n_col_domain = int(n_col_domain or cfg["data"]["n_col_domain"])
+    n_train_per_bc = int(n_train_per_bc or cfg["data"]["n_train_per_bc"])
+    dtype = dtype_from_config(cfg["model"].get("dtype", "float32"))
+
+    set_seed(int(cfg["seed"]))
+    x_col, x_train, sol_train = get_data_fluid(
+        problem=problem,
+        N_col_domain=n_col_domain,
+        N_train=n_train_per_bc,
+    )
+    collocation_x = x_col.to(device=device, dtype=dtype).clone().requires_grad_(True)
+
+    models = []
+    for i, name in enumerate(cfg["output_names"]):
+        model = GPPLUS(
+            train_x=x_train[i].type(dtype),
+            train_y=sol_train[i].type(dtype),
+            collocation_x=collocation_x,
+            basis=cfg["model"]["mean_function"],
+            NN_layers_base=cfg["model"]["nn_layers_base"],
+            name_output=name,
+            device=device,
+            dtype=dtype,
+        ).to(device=device, dtype=dtype)
+        models.append(model)
+
+    return models, {
+        "problem": problem,
+        "n_col_domain": n_col_domain,
+        "n_train_per_bc": n_train_per_bc,
+        "x_col_shape": tuple(x_col.shape),
+        "x_train_shapes": [tuple(x.shape) for x in x_train],
+        "sol_train_shapes": [tuple(y.shape) for y in sol_train],
+    }
+
+
+def save_checkpoint(path, model_list, cfg, metadata, loss_history):
+    path = resolve_path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "model_state_dicts": [model.state_dict() for model in model_list],
+            "config": cfg,
+            "metadata": metadata,
+            "loss_history": loss_history,
+        },
+        path,
+    )
+    return path
+
+
+def load_checkpoint(path, model_list, device):
+    path = resolve_path(path)
+    if not path.is_file():
+        raise FileNotFoundError(f"Missing checkpoint: {path}. Run scripts/train.py first.")
+    try:
+        checkpoint = torch.load(path, map_location=device, weights_only=False)
+    except TypeError:
+        checkpoint = torch.load(path, map_location=device)
+
+    state_dicts = checkpoint.get("model_state_dicts")
+    if state_dicts is None:
+        raise KeyError(f"Checkpoint does not contain model_state_dicts: {path}")
+    for model, state_dict in zip(model_list, state_dicts):
+        model.load_state_dict(state_dict)
+    return checkpoint
+
+
+def dump_json(data, path):
+    path = resolve_path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)
+    return path
+
+
+def tensor_to_numpy_dict(fields):
+    return {key: value.detach().cpu().numpy() for key, value in fields.items()}

scripts/fake_data.py

@@ -0,0 +1,65 @@
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from scripts.common import PROBLEMS, ensure_onescience_path, load_config, resolve_path
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Generate GP_for_TO runtime sample tensors.")
+    parser.add_argument("--problem", choices=PROBLEMS, default=None)
+    parser.add_argument("--n-col-domain", type=int, default=None)
+    parser.add_argument("--n-train-per-bc", type=int, default=None)
+    parser.add_argument("--output-dir", default=None)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    cfg = load_config()
+    ensure_onescience_path(cfg.get("runtime", {}).get("onescience_src"))
+    from onescience.utils.GP_TO import get_data_fluid, set_seed
+
+    problem = args.problem or cfg["fake_data"]["problem"]
+    n_col_domain = args.n_col_domain or cfg["fake_data"]["n_col_domain"]
+    n_train_per_bc = args.n_train_per_bc or cfg["fake_data"]["n_train_per_bc"]
+    output_dir = resolve_path(args.output_dir or cfg["fake_data"]["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    set_seed(int(cfg["seed"]))
+    x_col, x_train, sol_train = get_data_fluid(
+        problem=problem,
+        N_col_domain=n_col_domain,
+        N_train=n_train_per_bc,
+    )
+
+    npz_path = output_dir / f"{problem}_samples.npz"
+    arrays = {"x_col": x_col.cpu().numpy()}
+    for i, name in enumerate(cfg["output_names"]):
+        arrays[f"x_train_{name}"] = x_train[i].cpu().numpy()
+        arrays[f"target_{name}"] = sol_train[i].cpu().numpy()
+    np.savez(npz_path, **arrays)
+
+    metadata = {
+        "problem": problem,
+        "n_col_domain_requested": int(n_col_domain),
+        "n_train_per_bc": int(n_train_per_bc),
+        "x_col_shape": list(x_col.shape),
+        "x_train_shapes": [list(x.shape) for x in x_train],
+        "target_shapes": [list(y.shape) for y in sol_train],
+    }
+    metadata_path = output_dir / f"{problem}_metadata.json"
+    metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
+    print(f"Fake GP_for_TO tensors written to {npz_path}")
+    print(json.dumps(metadata, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,94 @@
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from scripts.common import (
+    PROBLEMS,
+    build_models,
+    ensure_onescience_path,
+    load_checkpoint,
+    load_config,
+    resolve_path,
+    select_device,
+    tensor_to_numpy_dict,
+)
+
+ensure_onescience_path()
+from scripts.topology_optimization import clear_cached_kernels, predict_fields, share_mean_module
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run GP_for_TO field inference from a checkpoint.")
+    parser.add_argument("--problem", choices=PROBLEMS, default=None)
+    parser.add_argument("--gpu", type=int, default=None)
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--n-col-domain", type=int, default=None)
+    parser.add_argument("--n-train-per-bc", type=int, default=None)
+    parser.add_argument("--checkpoint-path", default=None)
+    parser.add_argument("--output-dir", default=None)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    cfg = load_config()
+    if args.problem:
+        cfg["problem"] = args.problem
+    if args.gpu is not None:
+        cfg["inference"]["gpu"] = args.gpu
+    if args.device:
+        cfg["inference"]["device"] = args.device
+    if args.n_col_domain is not None:
+        cfg["inference"]["n_col_domain"] = args.n_col_domain
+    if args.n_train_per_bc is not None:
+        cfg["data"]["n_train_per_bc"] = args.n_train_per_bc
+    if args.checkpoint_path:
+        cfg["inference"]["checkpoint_path"] = args.checkpoint_path
+    if args.output_dir:
+        cfg["inference"]["output_dir"] = args.output_dir
+
+    os.chdir(PROJECT_ROOT)
+    device = select_device(cfg["inference"])
+    models, metadata = build_models(
+        cfg,
+        device,
+        n_col_domain=cfg["inference"].get("n_col_domain", cfg["data"]["n_col_domain"]),
+        n_train_per_bc=cfg["data"]["n_train_per_bc"],
+        problem=cfg["problem"],
+    )
+    checkpoint = load_checkpoint(cfg["inference"]["checkpoint_path"], models, device)
+    share_mean_module(models)
+    for model in models:
+        model.eval()
+    clear_cached_kernels(models)
+
+    with torch.no_grad():
+        fields = predict_fields(models)
+
+    output_dir = resolve_path(cfg["inference"]["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    npz_path = output_dir / "predictions.npz"
+    np.savez(npz_path, **tensor_to_numpy_dict(fields))
+
+    summary = {
+        **metadata,
+        "checkpoint_metadata": checkpoint.get("metadata", {}),
+        "output_file": str(npz_path),
+        "field_shapes": {key: list(value.shape) for key, value in fields.items()},
+    }
+    (output_dir / "inference_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    print(f"Saved predictions to {npz_path}")
+    print(json.dumps(summary["field_shapes"], indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,68 @@
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from scripts.common import load_config, resolve_path
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Summarize GP_for_TO inference outputs.")
+    parser.add_argument("--output-dir", default=None)
+    parser.add_argument("--no-plot", action="store_true")
+    return parser.parse_args()
+
+
+def save_plot(data, output_dir):
+    x = data["x"]
+    fields = [("u", data["u"]), ("v", data["v"]), ("p", data["p"]), ("ro", data["ro"])]
+    fig, axes = plt.subplots(2, 2, figsize=(9, 7))
+    for ax, (name, values) in zip(axes.reshape(-1), fields):
+        im = ax.tricontourf(x[:, 0], x[:, 1], values, levels=32, cmap="viridis")
+        ax.set_title(name)
+        ax.set_xlabel("x")
+        ax.set_ylabel("y")
+        fig.colorbar(im, ax=ax)
+    fig.tight_layout()
+    path = output_dir / "field_summary.png"
+    fig.savefig(path, dpi=160)
+    plt.close(fig)
+    return path
+
+
+def main():
+    args = parse_args()
+    cfg = load_config()
+    output_dir = resolve_path(args.output_dir or cfg["inference"]["output_dir"])
+    pred_path = output_dir / "predictions.npz"
+    summary_path = output_dir / "inference_summary.json"
+    if not pred_path.is_file():
+        raise FileNotFoundError(f"Missing inference output: {pred_path}")
+
+    data = np.load(pred_path)
+    print(f"Prediction file: {pred_path}")
+    for name in ("x", "u", "v", "p", "ro"):
+        arr = data[name]
+        print(
+            f"{name}: shape={arr.shape}, dtype={arr.dtype}, "
+            f"min={float(arr.min()):.6e}, max={float(arr.max()):.6e}, mean={float(arr.mean()):.6e}"
+        )
+
+    if summary_path.is_file():
+        summary = json.loads(summary_path.read_text(encoding="utf-8"))
+        print(f"Problem: {summary.get('problem')}, checkpoint source: {summary.get('checkpoint_metadata', {}).get('problem')}")
+
+    if not args.no_plot:
+        plot_path = save_plot(data, output_dir)
+        print(f"Plot: {plot_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/topology_optimization.py

@@ -0,0 +1,427 @@
+import torch
+import math
+import numpy as np
+from gpytorch.settings import cholesky_jitter
+import matplotlib.pyplot as plt
+import torch.optim
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+from onescience.utils.GP_TO import  plot_predictions_and_residuals,plot_loss_history,plot_density_and_velocity_fields
+plt.rcParams['font.family'] = 'DejaVu Sans'
+plt.rcParams['font.size'] = 15
+plt.rcParams['figure.dpi'] = 150
+
+
+checkpoints=[1000,10000,20000,30000,40000,50000]
+lambdaa = 0.1
+
+# Determine volume loss based on title
+gamma_values = {
+    'rugby': 0.9,
+    'pipebend': 0.08 * torch.pi,
+    'doublepipe': 1 / 3,
+    'diffuser': 0.5
+    }
+
+def modified_sigmoid(x, alpha=12.0):
+    """Compute the modified sigmoid function."""
+    return 1 / (1 + torch.exp(-alpha * (x - 0.5)))
+
+
+def infer_square_grid(num_points):
+    side = int(math.sqrt(num_points))
+    if side * side != num_points:
+        raise ValueError(
+            "Numerical differentiation expects a square collocation grid. "
+            f"Got {num_points} points."
+        )
+    return side, side
+
+
+def share_mean_module(model_list):
+    model_ref = model_list[0]
+    for model in model_list:
+        model.mean_module_NN_All = model_ref.mean_module_NN_All
+    return model_ref
+
+
+def clear_cached_kernels(model_list):
+    for model in model_list:
+        model.k_xX = None
+        model.chol_decomp = None
+
+
+def predict_fields(model_list):
+    model_ref = share_mean_module(model_list)
+    collocation_x = model_ref.collocation_x
+    m_col = model_ref.mean_module_NN_All(collocation_x)
+
+    if model_ref.k_xX is None:
+        for model in model_list:
+            model.k_xX = model.covar_module(model.train_inputs[0], collocation_x).evaluate()
+
+    if model_ref.chol_decomp is None:
+        with cholesky_jitter(1e-7):
+            for model in model_list:
+                model.chol_decomp = model.covar_module(model.train_inputs[0]).cholesky()
+
+    offsets = []
+    for i, model in enumerate(model_list):
+        target_offset = model.train_targets.unsqueeze(-1) - model.mean_module_NN_All(model.train_inputs[0])[:, i].unsqueeze(-1)
+        offsets.append(model.chol_decomp._cholesky_solve(target_offset))
+
+    u = (m_col[:, 0].unsqueeze(-1) + model_list[0].k_xX.t() @ offsets[0]).squeeze(-1)
+    v = (m_col[:, 1].unsqueeze(-1) + model_list[1].k_xX.t() @ offsets[1]).squeeze(-1)
+    p = (m_col[:, 2].unsqueeze(-1) + model_list[2].k_xX.t() @ offsets[2]).squeeze(-1)
+    ro_tensor = (m_col[:, 3].unsqueeze(-1) + model_list[3].k_xX.t() @ offsets[3]).squeeze(-1)
+    ro = modified_sigmoid(ro_tensor)
+    return {"x": collocation_x, "u": u, "v": v, "p": p, "ro": ro}
+
+
+def compute_dynamic_weights_ic(pde_loss, bc_loss, ic_loss, model):
+    """
+    Compute dynamic weights for loss functions based on gradients.
+    """
+    params_to_update = [param for param in model.parameters() if param.requires_grad]
+    
+    def compute_gradients(loss, scaling_factor):
+        gradients = torch.autograd.grad(scaling_factor * loss, params_to_update, retain_graph=True, allow_unused=True)
+        values = [p.reshape(-1).cpu().tolist() for p in gradients if p is not None]
+        return torch.abs(torch.tensor([v for val in values for v in val]))
+    
+    delta_pde = compute_gradients(pde_loss, 1.0)
+    delta_bc = compute_gradients(bc_loss, model.alpha)
+    delta_ic = compute_gradients(ic_loss, model.beta)
+
+    temp_bc = torch.max(delta_pde) / torch.mean(delta_bc)
+    temp_ic = torch.max(delta_pde) / torch.mean(delta_ic)
+
+    return (
+        (1.0 - lambdaa) * model.alpha + lambdaa * temp_bc,
+        (1.0 - lambdaa) * model.beta + lambdaa * temp_ic
+    )
+
+def compute_fvm_residuals_higher_order(u, v, p, f_x, f_y, nu, dx, dy):
+    # Check the device of the input tensors and move other tensors to the same device
+    device = u.device
+    
+    # Extend the fields with ghost cells to handle boundary conditions
+    u_ext = torch.zeros((u.shape[0] + 4, u.shape[1] + 4), device=device)
+    v_ext = torch.zeros((v.shape[0] + 4, v.shape[1] + 4), device=device)
+    p_ext = torch.zeros((p.shape[0] + 4, p.shape[1] + 4), device=device)
+    
+    # Copy the interior
+    u_ext[2:-2, 2:-2] = u
+    v_ext[2:-2, 2:-2] = v
+    p_ext[2:-2, 2:-2] = p
+    
+    # Apply ghost cells for no-slip boundary conditions
+    u_ext[:2, :] = u_ext[2:4, :]
+    u_ext[-2:, :] = u_ext[-4:-2, :]
+    u_ext[:, :2] = u_ext[:, 2:4]
+    u_ext[:, -2:] = u_ext[:, -4:-2]
+    
+    v_ext[:2, :] = v_ext[2:4, :]
+    v_ext[-2:, :] = v_ext[-4:-2, :]
+    v_ext[:, :2] = v_ext[:, 2:4]
+    v_ext[:, -2:] = v_ext[:, -4:-2]
+    
+    # Laplacian of u and v (fourth-order central differences for second derivatives)
+    u_xx = (-u_ext[4:, 2:-2] + 16 * u_ext[3:-1, 2:-2] - 30 * u_ext[2:-2, 2:-2] + 16 * u_ext[1:-3, 2:-2] - u_ext[:-4, 2:-2]) / (12 * dx**2)
+    u_yy = (-u_ext[2:-2, 4:] + 16 * u_ext[2:-2, 3:-1] - 30 * u_ext[2:-2, 2:-2] + 16 * u_ext[2:-2, 1:-3] - u_ext[2:-2, :-4]) / (12 * dy**2)
+    laplacian_u = u_xx + u_yy
+    
+    v_xx = (-v_ext[4:, 2:-2] + 16 * v_ext[3:-1, 2:-2] - 30 * v_ext[2:-2, 2:-2] + 16 * v_ext[1:-3, 2:-2] - v_ext[:-4, 2:-2]) / (12 * dx**2)
+    v_yy = (-v_ext[2:-2, 4:] + 16 * v_ext[2:-2, 3:-1] - 30 * v_ext[2:-2, 2:-2] + 16 * v_ext[2:-2, 1:-3] - v_ext[2:-2, :-4]) / (12 * dy**2)
+    laplacian_v = v_xx + v_yy
+    
+    # Gradients of p using fourth-order central difference
+    grad_p_x = (-p_ext[4:, 2:-2] + 8 * p_ext[3:-1, 2:-2] - 8 * p_ext[1:-3, 2:-2] + p_ext[:-4, 2:-2]) / (12 * dx)
+    grad_p_y = (-p_ext[2:-2, 4:] + 8 * p_ext[2:-2, 3:-1] - 8 * p_ext[2:-2, 1:-3] + p_ext[2:-2, :-4]) / (12 * dy)
+    
+    # Momentum equation residuals
+    residual_u = -nu * laplacian_u + grad_p_x + f_x
+    residual_v = -nu * laplacian_v + grad_p_y + f_y
+    
+    # Mass conservation (divergence of velocity) using central differences
+    div_u = (u_ext[2:-2, 2:-2] - u_ext[1:-3, 2:-2]) / dx
+    div_v = (v_ext[2:-2, 2:-2] - v_ext[2:-2, 1:-3]) / dy
+    
+    u_y = (u_ext[2:-2, 2:-2] - u_ext[2:-2, 1:-3]) / dy
+    v_x = (v_ext[2:-2, 2:-2] - v_ext[1:-3, 2:-2]) / dx
+    residual_mass = div_u + div_v
+    
+    return residual_u.reshape(-1), residual_v.reshape(-1), residual_mass.reshape(-1), div_u.reshape(-1), div_v.reshape(-1), u_y.reshape(-1), v_x.reshape(-1)
+
+
+
+def compute_autograd_derivatives(model, z_column, collocation_x, grad_order=1):
+    """
+    Compute first and second-order derivatives using PyTorch autograd.
+    :param model: Model to train
+    :param z_column: Column of `z_all` for differentiation
+    :param collocation_x: Input coordinates
+    :param grad_order: Order of gradients (1 for first, 2 for second)
+    :return: First and second-order derivatives as tensors
+    """
+    model.train()
+    grad_1 = torch.autograd.grad(z_column, collocation_x, torch.ones_like(z_column), create_graph=True)[0]
+    if grad_order > 1:
+        grad_2_x = torch.autograd.grad(grad_1[:, 0], collocation_x, torch.ones_like(grad_1[:, 0]), create_graph=True)[0][:, 0]
+        grad_2_y = torch.autograd.grad(grad_1[:, 1], collocation_x, torch.ones_like(grad_1[:, 1]), create_graph=True)[0][:, 1]
+        return grad_1[:, 0], grad_1[:, 1], grad_2_x, grad_2_y
+    return grad_1[:, 0], grad_1[:, 1], None, None
+
+w_1,w_2,w_3, w_4, w_5 =0.01,0.01,1e2,1e5,1e4
+
+def loss_volume(y, gamma=0.5):
+    """
+    Compute the volume loss as the squared difference between the mean of y and gamma.
+    """
+    mean_y = torch.mean(y)
+    return torch.square(mean_y - gamma)
+
+
+def dissipated_power(u, u_x, v_y, u_y, v_x):
+    """
+    Compute the total dissipated power based on the input velocity gradients and displacements.
+    """
+    # First part of the dissipated power
+    p1 = (u_x**2 + v_y**2 + u_y**2 + v_x**2).sum(dim=1, keepdim=True)
+    
+    # Second part of the dissipated power
+    u2 = (u[:, :2]**2).sum(dim=1, keepdim=True)
+    p2 = alpha(u[:, 3:]) * u2
+    
+    return 0.5 * (p1 + p2)
+
+
+def alpha(rho):
+    """
+    Compute the alpha parameter based on rho.
+    """
+    alpha_max = 2.5e4
+    alpha_min = 2.5e-4
+    q = 0.1
+    return alpha_max + (alpha_min - alpha_max) * rho * (1 + q) / (rho + q)
+
+def calculate_loss_multioutput(
+    model_list,
+    iteration,
+    diff_method='Numerical',
+    title="default",
+    problem="doublepipe",
+    checkpoint_steps=None,
+    plot_outputs=True,
+):
+    """
+    Calculate loss for multi-output models based on PDE, BC, and IC residuals.
+    """
+    # Clone collocation points
+    collocation_x = model_list[0].collocation_x.clone()
+
+    # Compute mean values at collocation points
+    m_col = model_list[0].mean_module_NN_All(collocation_x)
+
+    # Evaluate covariance matrices if not already done
+    if model_list[0].k_xX is None:
+        for i, model in enumerate(model_list):
+            model.k_xX = model.covar_module(model.train_inputs[0], collocation_x).evaluate()
+
+    # Perform Cholesky decomposition if not already done
+    if model_list[0].chol_decomp is None:
+        with cholesky_jitter(1e-7):
+            for model in model_list:
+                model.chol_decomp = model.covar_module(model.train_inputs[0]).cholesky()
+
+    # Solve for offsets using Cholesky decomposition
+    offsets = []
+    for i, model in enumerate(model_list):
+        target_offset = model.train_targets.unsqueeze(-1) - model.mean_module_NN_All(model.train_inputs[0])[:, i].unsqueeze(-1)
+        offsets.append(model.chol_decomp._cholesky_solve(target_offset))
+
+    # Compute values for velocity, pressure, and density
+    ro_tensor = (m_col[:, 3].unsqueeze(-1) + model_list[3].k_xX.t() @ offsets[3]).squeeze(-1)
+    ro = modified_sigmoid(ro_tensor)
+
+    u = (m_col[:, 0].unsqueeze(-1) + model_list[0].k_xX.t() @ offsets[0]).squeeze(-1)
+    v = (m_col[:, 1].unsqueeze(-1) + model_list[1].k_xX.t() @ offsets[1]).squeeze(-1)
+    p = (m_col[:, 2].unsqueeze(-1) + model_list[2].k_xX.t() @ offsets[2]).squeeze(-1)
+
+    z_all = torch.cat((u.unsqueeze(1), v.unsqueeze(1), p.unsqueeze(1), ro.unsqueeze(1)), dim=1)
+    f = alpha(z_all[:, 3:]) * z_all[:, :2]
+    fx, fy = f[:, :1], f[:, 1:]
+
+    # Compute residuals
+    if diff_method == 'Numerical':
+        nx, ny = infer_square_grid(collocation_x.shape[0])
+        dx, dy = 1.0 / nx, 1.0 / ny
+        u, v, p = [arr.reshape(nx, ny) for arr in (u, v, p)]
+        fx, fy = [arr.reshape(nx, ny) for arr in (fx, fy)]
+        residuals = compute_fvm_residuals_higher_order(u, v, p, fx, fy, nu=1, dx=dx, dy=dy)
+        residual_pde1, residual_pde2, residual_pde3, u_x, v_y, u_y, v_x = residuals
+        residual_pde1, residual_pde2, residual_pde3 = [res * w for res, w in zip(residuals[:3], (w_1, w_2, w_3))]
+        u, v, p = [arr.reshape(-1) for arr in (u, v, p)]
+        fx, fy = [arr.reshape(-1) for arr in (fx, fy)]
+    else:
+        u_x, u_y, u_xx, u_yy = compute_autograd_derivatives(model_list[0], z_all[:, 0], collocation_x, grad_order=2)
+        v_x, v_y, v_xx, v_yy = compute_autograd_derivatives(model_list[1], z_all[:, 1], collocation_x, grad_order=2)
+        p_x, p_y, _, _ = compute_autograd_derivatives(model_list[2], z_all[:, 2], collocation_x, grad_order=1)
+        residual_pde1 = (- (u_xx + u_yy) + p_x + fx[:, 0]) * w_1
+        residual_pde2 = (- (v_xx + v_yy) + p_y + fy[:, 0]) * w_2
+        residual_pde3 = (u_x + v_y) * w_3
+
+    # Plot fields at checkpoints
+    active_checkpoints = checkpoints if checkpoint_steps is None else checkpoint_steps
+    if plot_outputs and iteration in active_checkpoints:
+        plot_density_and_velocity_fields(u, v, p, ro, collocation_x, iteration, problem)
+        plot_predictions_and_residuals(u, v, p, ro, collocation_x, iteration, residual_pde1, residual_pde2, residual_pde3, 1, 1, 1, problem)
+    # Calculate losses
+    loss_pde1 = torch.mean(residual_pde1**2)
+    loss_pde2 = torch.mean(residual_pde2**2)
+    loss_pde3 = torch.mean(residual_pde3**2)
+    dp_loss = dissipated_power(z_all, u_x.reshape(-1, 1), v_y.reshape(-1, 1), u_y.reshape(-1, 1), v_x.reshape(-1, 1))
+
+    gamma = next((val for key, val in gamma_values.items() if key in title), 0.5)
+    vol_loss = loss_volume(z_all[:, 3:], gamma=gamma)* w_4
+
+    # Return final losses
+    return loss_pde1, loss_pde2, loss_pde3, torch.sum(dp_loss), vol_loss 
+
+
+
+
+def find_TO(
+    model_list,
+    lr_default: float = 0.01,
+    num_iter: int = 500,
+    title: str = 'default',
+    problem: str = 'doublepipe',   # 设置默认值
+    diff_method: str = 'Numerical',
+    checkpoint_steps=None,
+    plot_outputs=True,
+) -> float:
+    """
+    Train models to minimize total loss using dynamic weights and record progress.
+
+    Args:
+        model_list: List of models to be trained.
+        lr_default: Default learning rate for the optimizer.
+        num_iter: Number of training iterations.
+        title: Title for specific settings.
+        diff_method: Differentiation method ('Numerical' or 'Autograd').
+
+    Returns:
+        loss history
+    """
+    # Use the first model as a reference for shared NN
+    model_ref = share_mean_module(model_list)
+    for model in model_list:
+        model.train()
+
+    # Initialize variables
+    loss_total, loss_hist, loss_hist_total = [], [], []
+    loss_pde1_hist, loss_pde2_hist, loss_pde3_hist = [], [], []
+    loss_dp_hist, vol_loss_hist = [], []
+    scaled_loss_pde1_hist, scaled_loss_pde2_hist = [], []
+    scaled_loss_pde3_hist, scaled_loss_dp_hist, scaled_vol_loss_hist = [], [], []
+    weights = {'alpha': [], 'beta': [], 'mu_p': []}
+    result_loss_thetas = {'largest_eigval_hist': [], 'condition_number_hist': [], 'Hessian': [], 'grads': []}
+    GP_NN_hist, NN_hist, time_hist = [], [], []
+    dynamic_weights = True
+    sigma_1, beta_1, alph_1 = 10, 10, 0.1
+    f_inc, mu_F = math.inf, 1
+
+    # Initialize model-specific parameters
+    model1 = model_list[0]
+    model1.alpha, model1.beta, model1.theta = 1, 1, 1
+
+    # Set up optimizer and scheduler
+    optimizer = torch.optim.Adam(model_ref.parameters(), lr=lr_default)
+    milestones = sorted({int(x) for x in np.linspace(0, max(num_iter, 1), 4).tolist() if int(x) > 0})
+    scheduler = torch.optim.lr_scheduler.MultiStepLR(
+        optimizer,
+        milestones=milestones,
+        gamma=0.75
+    )
+
+    # Training loop
+    with tqdm(range(num_iter + 1), desc='Epoch', position=0, leave=True) as pbar:
+        for j in pbar:
+            optimizer.zero_grad()
+
+            # Calculate losses
+            loss_pde1, loss_pde2, loss_pde3, loss_dp, vol_loss = calculate_loss_multioutput(
+                model_list,
+                j,
+                diff_method=diff_method,
+                title=title,
+                problem=problem,
+                checkpoint_steps=checkpoint_steps,
+                plot_outputs=plot_outputs,
+            )
+            loss_pde = loss_pde1 + loss_pde3
+
+            # Dynamic weight adjustments
+            if dynamic_weights:
+                alpha, beta = compute_dynamic_weights_ic(
+                    loss_pde1 + loss_pde2, loss_pde3, vol_loss, model1
+                )
+                if all(torch.is_tensor(val) and not (torch.isnan(val).any() or torch.isinf(val).any()) for val in [alpha, beta]):
+                    model1.alpha, model1.beta = alpha, beta
+
+                weights['alpha'].append(alpha.detach().cpu().item())
+                weights['beta'].append(beta.detach().cpu().item())
+                weights['mu_p'].append(mu_F)
+
+                loss = loss_dp + mu_F * (
+                    loss_pde1 + loss_pde2 + model1.alpha * loss_pde3 + model1.beta * vol_loss
+                )
+
+                if (j + 1) % 50 == 0:
+                    mu_F = min(mu_F * 1.05, 5e2)
+            else:
+                alph_1 = 1e-6 / (j + 1)
+                beta_1, sigma_1 = 2 * (j + 1), 5 * (j + 1)
+                loss = sigma_1 * loss_pde + alph_1 * loss_dp + beta_1 * vol_loss
+
+            # Record loss history
+            loss_total.append(loss.item())
+            loss_pde1_hist.append((1/w_1) * loss_pde1.item())
+            loss_pde2_hist.append((1/w_2) * loss_pde2.item())
+            loss_pde3_hist.append((1/w_3) * loss_pde3.item())
+            loss_dp_hist.append((1/w_5) * loss_dp.item())
+            vol_loss_hist.append((1/w_4) * vol_loss.item())
+            scaled_loss_pde1_hist.append(loss_pde1.item())
+            scaled_loss_pde2_hist.append(loss_pde2.item())
+            scaled_loss_pde3_hist.append(model1.alpha * loss_pde3.item())
+            scaled_loss_dp_hist.append(loss_dp.item())
+            scaled_vol_loss_hist.append(model1.beta * vol_loss.item())
+
+            # Plot loss history at checkpoints
+            active_checkpoints = checkpoints if checkpoint_steps is None else checkpoint_steps
+            if plot_outputs and j in active_checkpoints:
+                plot_loss_history(
+                    loss_total,
+                    [loss_pde1_hist, scaled_loss_pde1_hist],
+                    [loss_pde2_hist, scaled_loss_pde2_hist],
+                    [loss_pde3_hist, scaled_loss_pde3_hist],
+                    [loss_dp_hist, scaled_loss_dp_hist],
+                    [vol_loss_hist, scaled_vol_loss_hist],
+                    j,
+                    problem
+                )
+
+            # Backpropagation and optimizer step
+            loss.backward(retain_graph=True)
+            optimizer.step()
+            scheduler.step()
+
+            # Update progress description
+            pbar.set_postfix(loss=f"{loss.item():.6f}")
+
+            loss_hist.append(loss.item())
+
+    loss_hist_total = loss_hist
+
+    return loss_hist_total

scripts/train.py

@@ -0,0 +1,116 @@
+import argparse
+import os
+import sys
+import time
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from scripts.common import (
+    PROBLEMS,
+    build_models,
+    dump_json,
+    ensure_onescience_path,
+    load_config,
+    resolve_path,
+    save_checkpoint,
+    select_device,
+)
+
+ensure_onescience_path()
+from scripts.topology_optimization import find_TO
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train GP_for_TO topology optimization models.")
+    parser.add_argument("--problem", choices=PROBLEMS, default=None)
+    parser.add_argument("--gpu", type=int, default=None)
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--num-iter", type=int, default=None)
+    parser.add_argument("--n-col-domain", type=int, default=None)
+    parser.add_argument("--n-train-per-bc", type=int, default=None)
+    parser.add_argument("--diff-method", choices=("Numerical", "Autograd"), default=None)
+    parser.add_argument("--lr", type=float, default=None)
+    parser.add_argument("--checkpoint-path", default=None)
+    parser.add_argument("--output-dir", default=None)
+    parser.add_argument("--no-plot", action="store_true")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    cfg = load_config()
+
+    if args.problem:
+        cfg["problem"] = args.problem
+    if args.gpu is not None:
+        cfg["training"]["gpu"] = args.gpu
+    if args.device:
+        cfg["training"]["device"] = args.device
+    if args.num_iter is not None:
+        cfg["training"]["num_iter"] = args.num_iter
+    if args.n_col_domain is not None:
+        cfg["data"]["n_col_domain"] = args.n_col_domain
+    if args.n_train_per_bc is not None:
+        cfg["data"]["n_train_per_bc"] = args.n_train_per_bc
+    if args.diff_method:
+        cfg["training"]["diff_method"] = args.diff_method
+    if args.lr is not None:
+        cfg["training"]["lr_default"] = args.lr
+    if args.checkpoint_path:
+        cfg["training"]["checkpoint_path"] = args.checkpoint_path
+    if args.output_dir:
+        cfg["training"]["output_dir"] = args.output_dir
+    if args.no_plot:
+        cfg["training"]["plot_outputs"] = False
+
+    os.chdir(PROJECT_ROOT)
+    device = select_device(cfg["training"])
+    models, metadata = build_models(
+        cfg,
+        device,
+        n_col_domain=cfg["data"]["n_col_domain"],
+        n_train_per_bc=cfg["data"]["n_train_per_bc"],
+        problem=cfg["problem"],
+    )
+
+    title = f"seed{cfg['seed']}_{cfg['problem']}_{datetime.now().strftime('%B%d_%H-%M')}"
+    start_time = time.time()
+    loss_history = find_TO(
+        model_list=models,
+        num_iter=int(cfg["training"]["num_iter"]),
+        lr_default=float(cfg["training"]["lr_default"]),
+        title=title,
+        problem=cfg["problem"],
+        diff_method=cfg["training"]["diff_method"],
+        checkpoint_steps=list(cfg["training"].get("checkpoints", [])),
+        plot_outputs=bool(cfg["training"].get("plot_outputs", True)),
+    )
+    elapsed = time.time() - start_time
+
+    output_dir = resolve_path(cfg["training"]["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    np.save(output_dir / "loss_history.npy", np.asarray(loss_history, dtype=np.float64))
+    summary = {
+        **metadata,
+        "device": str(device),
+        "num_iter": int(cfg["training"]["num_iter"]),
+        "diff_method": cfg["training"]["diff_method"],
+        "elapsed_seconds": elapsed,
+        "final_loss": float(loss_history[-1]) if loss_history else None,
+    }
+    dump_json(summary, output_dir / "training_summary.json")
+    ckpt_path = save_checkpoint(cfg["training"]["checkpoint_path"], models, cfg, summary, loss_history)
+
+    print(f"Training finished in {elapsed:.2f}s")
+    print(f"Loss history: {output_dir / 'loss_history.npy'}")
+    print(f"Checkpoint: {ckpt_path}")
+
+
+if __name__ == "__main__":
+    main()

weight/.gitkeep

@@ -0,0 +1 @@
+