unocg 0.0.2__tar.gz → 0.0.4__tar.gz

{unocg-0.0.2/unocg.egg-info → unocg-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unocg
-Version: 0.0.2
+Version: 0.0.4
 Summary: UNO-CG
 Author-email: Julius Herb <herb@mib.uni-stuttgart.de>, Felix Fritzen <fritzen@mib.uni-stuttgart.de>
 License-Expression: MIT
@@ -20,9 +20,9 @@ Requires-Dist: torch==2.10.0
 Requires-Dist: tqdm~=4.67.1
 Requires-Dist: torchvision~=0.25.0
 Requires-Dist: pyvista[jupyter]~=0.46.5
+Requires-Dist: jupyterlab~=4.5.3
 Provides-Extra: dev
 Requires-Dist: jupytext~=1.16.4; extra == "dev"
-Requires-Dist: jupyterlab~=4.5.3; extra == "dev"
 Requires-Dist: pre_commit~=3.8.0; extra == "dev"
 Requires-Dist: sphinx~=7.4.7; extra == "dev"
 Requires-Dist: sphinx-rtd-theme~=2.0.0; extra == "dev"

{unocg-0.0.2 → unocg-0.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unocg"
-version = "0.0.2"
+version = "0.0.4"
 authors = [
   { name="Julius Herb", email="herb@mib.uni-stuttgart.de" },
   { name="Felix Fritzen", email="fritzen@mib.uni-stuttgart.de" },
@@ -23,19 +23,19 @@ dependencies = [
     "torch==2.10.0",
     "tqdm~=4.67.1",
     "torchvision~=0.25.0",
-    "pyvista[jupyter]~=0.46.5"
+    "pyvista[jupyter]~=0.46.5",
+    "jupyterlab~=4.5.3",
 ]
 
 [project.optional-dependencies]
 dev = [
     "jupytext~=1.16.4",
-    "jupyterlab~=4.5.3",
     "pre_commit~=3.8.0",
     "sphinx~=7.4.7",
     "sphinx-rtd-theme~=2.0.0",
-    'myst-parser~=4.0.0',
-    'nbsphinx~=0.9.5',
-    'sphinx-gallery~=0.17.1',
+    "myst-parser~=4.0.0",
+    "nbsphinx~=0.9.5",
+    "sphinx-gallery~=0.17.1",
 ]
 all = [
     "unocg[dev]"

unocg-0.0.4/unocg/training/losses/__init__.py

@@ -0,0 +1 @@
+from .base import Loss, WeightedLoss

unocg-0.0.4/unocg/training/losses/base.py

@@ -0,0 +1,102 @@
+"""
+Definition of loss function that allow for a physics-informed training of machine learning models
+"""
+from abc import ABC
+from typing import Union, Iterable, List, Optional
+
+# third-party packages
+import torch
+
+
+class Loss(ABC):
+    """
+    Abstract loss function
+    """
+
+    def __init__(self, n_dim: int = 2, reduction: str = "sum"):
+        """
+
+        :param n_dim: dimension of the problem
+        :param reduction: ('mean'|'sum'|'none'), defaults to 'sum'
+        :type reduction: str, optional
+        """
+        super().__init__()
+        self.n_dim = n_dim
+        self.ch_dim = -(1 + self.n_dim)
+        self.reduction = reduction
+
+    def reduce(self, loss):
+        """
+        Perform a reduction step over all datasets to transform a loss function to a cost function.
+
+        A loss function is evaluated element-wise for a dataset.
+        However, a cost function should return a single value for the dataset.
+        Typically, `mean` reduction is used.
+
+        :param loss: Tensor that contains the element-wise loss for a dataset
+        :type loss: :class:`torch.Tensor`
+        :return: Reduced loss
+        :rtype: float
+        """
+        if self.reduction == "mean":
+            return torch.nanmean(loss)
+        elif self.reduction == "sum":
+            return torch.nansum(loss)
+        else:
+            return loss
+
+    def unsqueeze(self, output, target):
+        """
+        Ensure that the tensors :code:`output` and :code:`target` have a shape of the form :code:`(N, features)`.
+
+        When a loss function is called with a single data point, the tensor shape is :code:`(features)` and hence does not fit.
+        This method expands the dimensions if needed.
+
+        :param output: Model output
+        :type output: :class:`torch.Tensor`
+        :param target: Target data
+        :type target: :class:`torch.Tensor`
+        :return: Tuple (output, target)
+        :rtype: tuple
+        """
+        while output.ndim < self.n_dim + 2:
+            output = torch.unsqueeze(output, 0)
+        while target.ndim < output.ndim:
+            target = torch.unsqueeze(target, 0)
+        while output.ndim < target.ndim:
+            output = torch.unsqueeze(output, 0)
+        assert output.ndim >= self.n_dim + 2 and output.ndim == target.ndim
+        return output, target
+        # return output.flatten(end_dim=-(self.n_dim + 2)), target.flatten(end_dim=-(self.n_dim + 2))
+
+    @property
+    def dims(self):
+        return self.n_dim * (slice(None),)
+
+    @property
+    def dims_list(self):
+        return tuple(range(-1, -(1 + self.n_dim), -1))
+
+    @property
+    def expand_dims(self):
+        return self.n_dim * (None,)
+
+
+class WeightedLoss(Loss):
+    """
+    Weighted loss function that represents a linear combination of several loss functions
+    """
+
+    def __init__(self, losses: Iterable[any], weights: Iterable[Union[float, int]], reduction="mean"):
+        super().__init__(reduction=reduction)
+        self.losses = losses
+        self.weights = weights
+
+    def __call__(self, output, target):
+        total_loss = 0.0
+        for loss, weight in zip(self.losses, self.weights):
+            total_loss += weight * loss(output, target)
+        return total_loss
+
+    def __str__(self):
+        return f"Weighted(...)"

unocg-0.0.4/unocg/training/losses/mechanical.py

@@ -0,0 +1,246 @@
+import torch
+from .base import Loss
+
+
+class DispLoss(Loss):
+    """
+    Displacement loss function
+    """
+
+    def __init__(self, n_dim=2, reduction="sum"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+
+    def abs(self, output, target):
+        """
+        Compute absolute error for temperature field
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_disp, target_disp = output[..., :self.n_dim, *self.dims], target[..., :self.n_dim, *self.dims]
+        loss = torch.sqrt(torch.nn.MSELoss(reduction="none")(target_disp, output_disp).nanmean(self.dims_list))
+        loss = self.reduce(loss).squeeze(-1)
+        return loss
+
+    def rel(self, output, target):
+        """
+        Compute relative error for temperature field
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_disp, target_disp = output[..., :self.n_dim, *self.dims], target[..., :self.n_dim, *self.dims]
+        loss = torch.linalg.norm(target_disp - output_disp, dim=-(1 + self.n_dim)).nanmean(self.dims_list) \
+            / torch.linalg.norm(target_disp, dim=-(1 + self.n_dim)).nanmean(self.dims_list)
+        loss = self.reduce(loss).squeeze(-1)
+        return loss
+
+    def __call__(self, x, y):
+        return self.rel(x, y)
+
+    def __str__(self):
+        return f"Temp({self.reduction})"
+
+
+class StrainLoss(Loss):
+    """
+    Strain loss function
+    """
+
+    def __init__(self, grad_module, n_dim=2, reduction="sum"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+        self.grad_module = grad_module
+
+    def abs(self, output, target):
+        """
+        Compute absolute error for temperature field
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :self.n_dim, *self.dims], target[..., :self.n_dim, *self.dims]
+        loss = torch.linalg.norm(target_temp - output_temp, dim=self.ch_dim)
+        loss = self.reduce(loss)
+        return loss
+
+    def rel(self, output, target):
+        """
+        Compute relative error for temperature field
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :self.n_dim, *self.dims], target[..., :self.n_dim, *self.dims]
+        batch_dims = output_temp.shape[:self.ch_dim]
+        output_temp = torch.flatten(output_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        target_temp = torch.flatten(target_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        output_grad = self.grad_module(output_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        target_grad = self.grad_module(target_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        output_temp = torch.unflatten(output_temp, dim=0, sizes=batch_dims)
+        target_temp = torch.unflatten(target_temp, dim=0, sizes=batch_dims)
+        loss = torch.linalg.norm(target_grad - output_grad, dim=-(1 + self.n_dim)).nanmean(self.dims_list)
+        loss = self.reduce(loss)
+        return loss
+
+    def __call__(self, x, y):
+        return self.rel(x, y)
+
+    def __str__(self):
+        return f"Grad({self.reduction})"
+
+
+class StressLoss(Loss):
+    """
+    Flux loss function
+    """
+
+    def __init__(self, n_dim=2, reduction="sum"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+
+    def abs(self, output, target):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_flux, target_flux = output[..., self.n_dim:, *self.dims], target[..., self.n_dim:, *self.dims]
+        loss = torch.linalg.norm(target_flux - output_flux, dim=-(1 + self.n_dim)).nanmean(self.dims_list)
+        loss = self.reduce(loss)
+        return loss
+
+    def rel(self, output, target):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_flux, target_flux = output[..., self.n_dim:, *self.dims], target[..., self.n_dim:, *self.dims]
+        loss = torch.linalg.norm(target_flux - output_flux, dim=-(1 + self.n_dim)).nanmean(self.dims_list) \
+            / torch.linalg.norm(target_flux, dim=-(1 + self.n_dim)).nanmean(self.dims_list)
+        loss = self.reduce(loss)
+        return loss
+
+    def __call__(self, output: torch.Tensor, target: torch.Tensor):
+        return self.rel(output, target)
+
+    def __str__(self):
+        return f"Flux({self.reduction})"
+    
+    
+class MechanicalEnergyLoss(Loss):
+    """
+    Thermal energy loss function
+    """
+
+    def __init__(self, grad_module, n_dim=2, reduction="sum"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+        self.grad_module = grad_module
+
+    def abs(self, output, target):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :self.n_dim, *self.dims], target[..., :self.n_dim, *self.dims]
+        output_flux, target_flux = output[..., self.n_dim:, *self.dims], target[..., self.n_dim:, *self.dims]
+        batch_dims = output_temp.shape[:self.ch_dim]
+        output_temp = torch.flatten(output_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        target_temp = torch.flatten(target_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        output_grad = self.grad_module(output_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        target_grad = self.grad_module(target_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        output_temp = torch.unflatten(output_temp, dim=0, sizes=batch_dims)
+        target_temp = torch.unflatten(target_temp, dim=0, sizes=batch_dims)
+        loss = torch.sqrt(
+            torch.linalg.norm((output_grad - target_grad) * (output_flux - target_flux), dim=self.ch_dim).nanmean(self.dims_list)
+        )
+        loss = self.reduce(loss)
+        return loss
+
+    def rel(self, output, target):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :self.n_dim, *self.dims], target[..., :self.n_dim, *self.dims]
+        output_flux, target_flux = output[..., self.n_dim:, *self.dims], target[..., self.n_dim:, *self.dims]
+        batch_dims = output_temp.shape[:self.ch_dim]
+        output_temp = torch.flatten(output_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        target_temp = torch.flatten(target_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        output_grad = self.grad_module(output_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        target_grad = self.grad_module(target_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        output_temp = torch.unflatten(output_temp, dim=0, sizes=batch_dims)
+        target_temp = torch.unflatten(target_temp, dim=0, sizes=batch_dims)
+        loss = torch.sqrt(
+            torch.linalg.norm((output_grad - target_grad) * (output_flux - target_flux), dim=self.ch_dim).nanmean(self.dims_list)
+            / torch.linalg.norm(target_grad * target_flux, dim=self.ch_dim).nanmean(self.dims_list)
+        )
+        loss = self.reduce(loss)
+        return loss
+
+    def __call__(self, output_disp: torch.Tensor, target_disp: torch.Tensor, param_field=None):
+        return self.abs(output_disp, target_disp)
+
+    def __str__(self):
+        return f"ThermEnergy({self.reduction})"
+
+
+class LinearElasticResidualLoss(Loss):
+    """
+    Homogenized stress loss function
+    """
+
+    def __init__(self, div_module, n_dim=2, reduction="sum", residual_mode="mean"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+        self.div_module = div_module
+        self.residual_mode = residual_mode
+
+    def abs(self, output, target=None):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_flux, target_flux = output[..., 1:, *self.dims], target[..., 1:, *self.dims]
+        batch_dims = target_flux.shape[:self.ch_dim]
+        output_flux = torch.flatten(output_flux, start_dim=0, end_dim=self.ch_dim - 1).unsqueeze(self.ch_dim - 1)
+        output_residual = self.div_module(output_flux)
+        output_residual = torch.unflatten(output_residual, dim=0, sizes=batch_dims)
+        if self.residual_mode == "sum":
+            loss = torch.abs(output_residual).nansum(self.dims_list)
+        elif self.residual_mode == "mean":
+            loss = torch.abs(output_residual).nanmean(self.dims_list)
+        else:
+            raise ValueError("Unknown residual_mode")
+        loss = self.reduce(loss).squeeze(-1)
+        return loss
+
+    def __call__(self, output: torch.Tensor, target):
+        return self.abs(output, target)
+
+    def __str__(self):
+        return f"HeatCondResidual({self.reduction})"

unocg-0.0.4/unocg/training/losses/thermal.py

@@ -0,0 +1,247 @@
+import torch
+from .base import Loss
+
+
+class TempLoss(Loss):
+    """
+    Displacement loss function
+    """
+
+    def __init__(self, n_dim=2, reduction="sum"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+
+    def abs(self, output, target):
+        """
+        Compute absolute error for temperature field
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :1, *self.dims], target[..., :1, *self.dims]
+        loss = torch.sqrt(torch.nn.MSELoss(reduction="none")(target_temp, output_temp).nanmean(self.dims_list))
+        loss = self.reduce(loss).squeeze(-1)
+        return loss
+
+    def rel(self, output, target):
+        """
+        Compute relative error for temperature field
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :1, *self.dims], target[..., :1, *self.dims]
+        loss = torch.sqrt(torch.nn.MSELoss(reduction="none")(target_temp, output_temp).nanmean(self.dims_list) / \
+            torch.nn.MSELoss(reduction="none")(target_temp, torch.zeros_like(target_temp)).nanmean(self.dims_list))
+        loss = self.reduce(loss).squeeze(-1)
+        return loss
+
+    def __call__(self, x, y):
+        return self.rel(x, y)
+
+    def __str__(self):
+        return f"Temp({self.reduction})"
+
+
+class GradLoss(Loss):
+    """
+    Displacement loss function
+    """
+
+    def __init__(self, grad_module, n_dim=2, reduction="sum"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+        self.grad_module = grad_module
+
+    def abs(self, output, target):
+        """
+        Compute absolute error for temperature field
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :1, *self.dims], target[..., :1, *self.dims]
+        loss = torch.linalg.norm(target_temp - output_temp, dim=-3)
+        loss = self.reduce(loss)
+        return loss
+
+    def rel(self, output, target):
+        """
+        Compute relative error for temperature field
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :1, *self.dims], target[..., :1, *self.dims]
+        # output_flux, target_flux = output[..., 1:, *self.dims], target[..., 1:, *self.dims]
+        batch_dims = output_temp.shape[:self.ch_dim]
+        output_temp = torch.flatten(output_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        target_temp = torch.flatten(target_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        output_grad = self.grad_module(output_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        target_grad = self.grad_module(target_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        output_temp = torch.unflatten(output_temp, dim=0, sizes=batch_dims)
+        target_temp = torch.unflatten(target_temp, dim=0, sizes=batch_dims)
+        loss = torch.linalg.norm(target_grad - output_grad, dim=-(1 + self.n_dim)).nanmean(self.dims_list)
+        loss = self.reduce(loss)
+        return loss
+
+    def __call__(self, x, y):
+        return self.rel(x, y)
+
+    def __str__(self):
+        return f"Grad({self.reduction})"
+
+
+class FluxLoss(Loss):
+    """
+    Flux loss function
+    """
+
+    def __init__(self, n_dim=2, reduction="sum"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+
+    def abs(self, output, target):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_flux, target_flux = output[..., 1:, *self.dims], target[..., 1:, *self.dims]
+        loss = torch.linalg.norm(target_flux - output_flux, dim=-(1 + self.n_dim)).nanmean(self.dims_list)
+        loss = self.reduce(loss)
+        return loss
+
+    def rel(self, output, target):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_flux, target_flux = output[..., 1:, *self.dims], target[..., 1:, *self.dims]
+        loss = torch.linalg.norm(target_flux - output_flux, dim=-(1 + self.n_dim)).nanmean(self.dims_list) \
+            / torch.linalg.norm(target_flux, dim=-(1 + self.n_dim)).nanmean(self.dims_list)
+        loss = self.reduce(loss)
+        return loss
+
+    def __call__(self, output: torch.Tensor, target: torch.Tensor):
+        return self.rel(output, target)
+
+    def __str__(self):
+        return f"Flux({self.reduction})"
+    
+    
+class ThermalEnergyLoss(Loss):
+    """
+    Thermal energy loss function
+    """
+
+    def __init__(self, grad_module, n_dim=2, reduction="sum"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+        self.grad_module = grad_module
+
+    def abs(self, output, target):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :1, *self.dims], target[..., :1, *self.dims]
+        output_flux, target_flux = output[..., 1:, *self.dims], target[..., 1:, *self.dims]
+        batch_dims = output_temp.shape[:self.ch_dim]
+        output_temp = torch.flatten(output_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        target_temp = torch.flatten(target_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        output_grad = self.grad_module(output_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        target_grad = self.grad_module(target_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        output_temp = torch.unflatten(output_temp, dim=0, sizes=batch_dims)
+        target_temp = torch.unflatten(target_temp, dim=0, sizes=batch_dims)
+        loss = torch.sqrt(
+            torch.linalg.norm((output_grad - target_grad) * (output_flux - target_flux), dim=self.ch_dim).nanmean(self.dims_list)
+        )
+        loss = self.reduce(loss)
+        return loss
+
+    def rel(self, output, target):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_temp, target_temp = output[..., :1, *self.dims], target[..., :1, *self.dims]
+        output_flux, target_flux = output[..., 1:, *self.dims], target[..., 1:, *self.dims]
+        batch_dims = output_temp.shape[:self.ch_dim]
+        output_temp = torch.flatten(output_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        target_temp = torch.flatten(target_temp, start_dim=0, end_dim=self.ch_dim - 1)
+        output_grad = self.grad_module(output_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        target_grad = self.grad_module(target_temp).nanmean(-(2 + self.n_dim))  # average over gauss points
+        output_temp = torch.unflatten(output_temp, dim=0, sizes=batch_dims)
+        target_temp = torch.unflatten(target_temp, dim=0, sizes=batch_dims)
+        loss = torch.sqrt(
+            torch.linalg.norm((output_grad - target_grad) * (output_flux - target_flux), dim=self.ch_dim).nanmean(self.dims_list)
+            / torch.linalg.norm(target_grad * target_flux, dim=self.ch_dim).nanmean(self.dims_list)
+        )
+        loss = self.reduce(loss)
+        return loss
+
+    def __call__(self, output_disp: torch.Tensor, target_disp: torch.Tensor, param_field=None):
+        return self.abs(output_disp, target_disp)
+
+    def __str__(self):
+        return f"ThermEnergy({self.reduction})"
+
+
+class HeatCondResidualLoss(Loss):
+    """
+    Homogenized stress loss function
+    """
+
+    def __init__(self, div_module, n_dim=2, reduction="sum", residual_mode="mean"):
+        super().__init__(n_dim=n_dim, reduction=reduction)
+        self.div_module = div_module
+        self.residual_mode = residual_mode
+
+    def abs(self, output, target=None):
+        """
+        Compute absolute error in strain norm
+
+        :param output:
+        :param target:
+        :return:
+        """
+        output, target = self.unsqueeze(output, target)
+        output_flux, target_flux = output[..., 1:, *self.dims], target[..., 1:, *self.dims]
+        batch_dims = target_flux.shape[:self.ch_dim]
+        output_flux = torch.flatten(output_flux, start_dim=0, end_dim=self.ch_dim - 1).unsqueeze(self.ch_dim - 1)
+        output_residual = self.div_module(output_flux)
+        output_residual = torch.unflatten(output_residual, dim=0, sizes=batch_dims)
+        if self.residual_mode == "sum":
+            loss = torch.abs(output_residual).nansum(self.dims_list)
+        elif self.residual_mode == "mean":
+            loss = torch.abs(output_residual).nanmean(self.dims_list)
+        else:
+            raise ValueError("Unknown residual_mode")
+        loss = self.reduce(loss).squeeze(-1)
+        return loss
+
+    def __call__(self, output: torch.Tensor, target):
+        return self.abs(output, target)
+
+    def __str__(self):
+        return f"HeatCondResidual({self.reduction})"

unocg-0.0.4/unocg/utils/data.py

@@ -0,0 +1,95 @@
+"""
+Data loading for microstructure datasets
+"""
+
+import math
+from typing import Optional, Union, Dict
+
+import torch
+import h5py
+
+
+class MicrostructureDataset(torch.utils.data.Dataset):
+    """
+    Represents a dataset in a microstructure from a HDF5 file
+    """
+    def __init__(
+        self,
+        file_name: str,
+        group_name: str,
+        lazy_loading: Optional[bool] = True,
+        device = "cpu",
+        dtype: Optional[torch.dtype] = None,
+    ):
+        """
+        Constructor of the class. Create a `PyTorch` dataset from given HDF5 file groups.
+
+        :param file_name: path to the HDF5 file
+        :type file_name: str
+        :param group_name: path to the group in the HDF5 file
+        :type group_name: str
+        :param lazy_loading:
+        :type lazy_loading: bool
+        :param dtype:
+        """
+        super().__init__()
+        self.file_name = file_name
+        self.group_name = group_name
+        self.lazy_loading = lazy_loading
+        self.keys = []
+        self.loaded_keys = []
+        self.images = {}
+        if dtype is None:
+            self.dtype = torch.float32
+        else:
+            self.dtype = dtype
+        self.device = device
+        self.tensor_args = {"dtype": self.dtype, "device": self.device}
+
+        with h5py.File(self.file_name, "r") as file:
+            for dset_name in file[self.group_name].keys():
+                self.keys.append(dset_name)
+        if not self.lazy_loading:
+            for dset_name in self.keys:
+                self.load_dset(dset_name)
+
+    def __len__(self) -> int:
+        """
+        Get the length of the dataset, i.e. how many data points it contains.
+
+        :return: Length of the dataset
+        :rtype: int
+        """
+        return len(self.keys)
+
+    def __getitem__(self, index: int) -> Dict[str, Union[torch.tensor, str]]:
+        """
+        Fetch a data point with given index from the dataset
+
+        :param index: Index of the data point
+        :type index: int
+        :return: microstructure image
+        :rtype: torch.Tensor
+        """
+        if index >= len(self.keys):
+            raise ValueError("Dataset is not available")
+        dset_name = self.keys[index]
+        self.load_dset(dset_name, force_loading=False)
+        return self.images[dset_name].clone()
+
+    def load_dset(self, dset_name: str, force_loading: Optional[bool] = True):
+        """
+        Load dataset from HDF5 file
+
+        :param dset_name:
+        :param force_loading:
+        :return:
+        """
+        if (not force_loading) and (dset_name in self.loaded_keys):
+            return
+        
+        with h5py.File(self.file_name, "r") as file:
+            image = torch.tensor(file[self.group_name][dset_name]["image"][...], **self.tensor_args)
+        
+        self.images[dset_name] = image
+        self.loaded_keys.append(dset_name)

unocg-0.0.4/unocg/utils/evaluation.py

@@ -0,0 +1,48 @@
+from typing import Dict
+
+import torch
+from unocg.problems import Problem
+from unocg.solvers.torch import CgSolver
+from unocg.modules.solvers import CgModule
+from unocg.preconditioners import Preconditioner
+import timeit
+
+
+def benchmark(model, inputs, device, n_runs=2, n_warmup=1, verbose=True):
+    def run_model():
+        with torch.inference_mode():
+            model(*inputs)
+            torch.cuda.synchronize(device)
+
+    for _ in range(n_warmup):
+        run_model()
+    model_time = timeit.timeit(run_model, number=n_runs) / n_runs
+    if verbose:
+        print(f"Runtime per execution: {model_time*1000.:.4f}ms")
+    return model_time
+
+
+def benchmark_cg(cg_module, param_fields, loadings, device, n_runs=2, n_warmup=1, verbose=True):
+    with torch.inference_mode():
+        guess = cg_module.zero_guess(param_fields, loadings)
+
+        inputs_module = (guess, param_fields, loadings)
+
+        print(f"Overall solver:")
+        u = cg_module(*inputs_module)
+        time_overall = benchmark(cg_module, inputs_module, device, n_runs=n_runs)
+
+        iter_layer = cg_module.iteration_layers[0]
+        print(f"Solver iteration:")
+        batch_shape = (*param_fields.shape[:(-cg_module.n_dim - 1)], loadings.shape[0])
+        field_shape = (*batch_shape, *cg_module.shape)
+        iv_fields, iv_scalars = iter_layer.init_internal_variables(batch_shape=batch_shape, init_residual=u)
+        time_iter = benchmark(iter_layer, (u, u, param_fields, loadings, iv_fields, iv_scalars), device, n_runs=n_runs*10)
+
+        print(f"Preconditioner application:")
+        time_prec = benchmark(cg_module.prec_model, (u,), n_runs=n_runs*100, n_warmup=n_warmup, device=device);
+
+        print(f"Residual computation:")
+        time_matvec = benchmark(cg_module.matvec_model, (u, param_fields), n_runs=n_runs*10, n_warmup=n_warmup, device=device)
+
+        return time_overall, time_iter, time_prec, time_matvec

unocg-0.0.4/unocg/utils/plotting.py

@@ -0,0 +1,339 @@
+"""
+Plotting utilities
+"""
+from typing import Optional, Iterable, List, Tuple
+
+import matplotlib
+from matplotlib.tri import Triangulation
+import numpy as np
+import torch
+from matplotlib import pyplot as plt, ticker
+import math
+import shutil
+from unocg.problems import Problem, BC
+
+try:
+    import pyvista as pv
+except:
+    pass
+
+plt.rcParams["text.usetex"] = True if shutil.which('latex') else False
+matplotlib.rcParams["text.latex.preamble"] = r"\usepackage{amsmath}"
+plot_width = 6.3
+
+colors = [
+    "g",
+    "#00BEFF",
+    "#004191",
+    "r",
+    "k"
+]
+
+
+def plot_ms(image: torch.Tensor, file: Optional[str] = None, show_axis: bool = False, show_cbar: bool = False):
+    """
+    Plot microstructure image
+
+    :param image:
+    :param file:
+    :param show_axis:
+    :param show_cbar:
+    :return:
+    """
+    image = image.detach().cpu()
+    fig, ax = plt.subplots(1, 1, figsize=[6.3 / 2, 2.65 if show_cbar else 6.3 / 2], dpi=300)
+    ms_cmap = plt.get_cmap("viridis", 2)
+    im = ax.imshow(image, origin="lower", interpolation="none", extent=(-0.5, 0.5, -0.5, 0.5), cmap=ms_cmap)
+
+    if show_cbar:
+        cb = fig.colorbar(im, ax=ax, ticks=[0, 1])
+        cb.ax.set_title(rf"$\chi_1$")
+        cb.ax.set_yticklabels(["$0$", "$1$"])
+
+    if show_axis:
+        ax.axis("on")
+        ax.set_xlim(-0.5, 0.5)
+        ax.set_ylim(-0.5, 0.5)
+        ax.set_xticks([-0.5, 0.0, 0.5])
+        ax.set_yticks([-0.5, 0.0, 0.5])
+        ax.set_xlabel(r"$\frac{x_1}{l_1}$")
+        ax.set_ylabel(r"$\frac{x_2}{l_2}$")
+    else:
+        ax.axis("off")
+
+    plt.tight_layout()
+    if file is not None:
+        plt.savefig(file, dpi=300)
+    plt.show()
+
+
+def plot_prec_action(ax, field_ref, init_res, precs, labels, ch_idx=0, plot_res=False, centered=True, ch_label=None):
+    idx = (0,0,ch_idx)
+    ax_i = 0
+    if centered:
+        s_abs_max = torch.max(field_ref[idx].cpu().abs())
+        s_max, s_min = s_abs_max, -s_abs_max
+    else:
+        s_max, s_min = torch.max(field_ref[idx].cpu()), torch.min(field_ref[idx].cpu())
+    
+    prec_actions = []
+    with torch.inference_mode():
+        for prec in precs:
+            prec_action = prec.apply_field(init_res)
+            prec_action = prec_action / prec_action[idx].max() * field_ref[idx].abs().max()
+            prec_actions.append(prec_action)
+            if centered:
+                s_abs_max = torch.maximum(s_abs_max, torch.max(prec_action[idx].cpu().abs()))
+                s_max, s_min = s_abs_max, -s_abs_max
+            else:
+                s_max = torch.maximum(s_max, torch.max(prec_action[idx].cpu()))
+                s_min = torch.minimum(s_min, torch.min(prec_action[idx].cpu()))
+
+    if plot_res:
+        im = ax[0,ax_i].imshow(init_res[idx].cpu().detach(), origin="lower", cmap="seismic")
+        if ch_label is None:
+            ax[0,ax_i].set_title(r"$\boldsymbol{r}^{(0)}$")
+        else:
+            ax[0,ax_i].set_title(rf"$\left( \boldsymbol{{r}}^{{(0)}} \right)_{{{ch_label}}}$")
+        plt.colorbar(im, ax=ax[ax_i])
+        ax_i += 1
+
+    im = ax[0,ax_i].imshow(field_ref[idx].cpu().detach(), origin="lower", cmap="jet", vmax=s_max, vmin=s_min)
+    if ch_label is None:
+        ax[0,ax_i].set_title(r"$\boldsymbol{u}_{\underline{\mu}}$")
+    else:
+        ax[0,ax_i].set_title(rf"$\left( \boldsymbol{{u}}_{{\underline{{a}}}} \right)_{{{ch_label}}}$")
+    plt.colorbar(im, ax=ax[0,ax_i])
+    ax_i += 1
+    
+    for prec_action, label in zip(prec_actions, labels):
+        im = ax[0,ax_i].imshow(prec_action[idx].cpu().detach(), origin="lower", cmap="jet", vmax=s_max, vmin=s_min)
+        if ch_label is None:
+            ax[0,ax_i].set_title(rf"$\boldsymbol{{s}}_\text{{{label}}}$")
+        else:
+            ax[0,ax_i].set_title(rf"$\left( \boldsymbol{{s}}_\text{{{label}}} \right)_{{{ch_label}}}$")
+        plt.colorbar(im, ax=ax[0,ax_i])
+        ax_i += 1
+
+    for ax_handle in ax.ravel():
+        ax_handle.set_xticks([])
+        ax_handle.set_yticks([])
+
+
+def plot_convergence(ax, results, labels, colors, metrics, metric_labels, load_names=None, rates=None, show_load_labels=False,
+                     xmin=0, xmax=None, bounds=False, ymin=1e-10, ymax=1e0, zoom=False, zoom_it = 10, zoom_tol = 1e-3):
+    """
+    
+    """
+    if load_names is None:
+        load_names = ["x"]
+    if rates is None:
+        rates = (None,) * len(results)
+    colors = colors[:len(results)]
+    while len(labels) < len(results):
+        labels.append("CG")
+
+    for load_i, load_name in enumerate(load_names):
+        if show_load_labels:
+            ax[0, load_i].set_title(f"loading {load_name}")
+
+        for metric_i, (metric, metric_label) in enumerate(zip(metrics, metric_labels)):
+            ax[metric_i, load_i].set_xlabel(r"iterations $[-]$")
+            ax[metric_i, load_i].set_ylabel(metric_label)
+            ax[metric_i, load_i].set_ylim(bottom=ymin, top=ymax)
+            if zoom:
+                axins = ax[metric_i, load_i].inset_axes([0.2, 0.15, 0.1, 0.45],
+                    xlim=(-1, zoom_it), ylim=(zoom_tol, 1e0), xticklabels=[], yticklabels=[])
+                axins.grid()
+                axins.set_xticks([0, zoom_it], ["$0$", f"${zoom_it}$"])
+
+            for result, label, color, rate in zip(results, labels, colors, rates):
+                if result is None:
+                    continue
+                iters = torch.arange(result["err_history"].shape[0])
+                metric_losses = result["losses"][metric]
+                if metric_losses.ndim == 1:
+                    metric_losses = metric_losses.unsqueeze(-1)
+                ax[metric_i, load_i].semilogy(iters, metric_losses[:, load_i].cpu() / metric_losses[0, load_i].cpu(), '-', c=color, label=label)
+
+                if zoom:
+                    axins.semilogy(iters, metric_losses[:, load_i].cpu() / metric_losses[0, load_i].cpu(), '-', c=color)
+
+            if zoom:
+                ax[metric_i, load_i].indicate_inset_zoom(axins, edgecolor="black")
+    
+    for ax_handle in ax.ravel():
+        ax_handle.grid()
+        ax_handle.set_xlim(left=xmin, right=xmax)
+
+
+def plot_convergence_histogram(ax, results, labels, colors, rtol=1e-6, bins=None, xmin=0, xmax=None, log_scale=False, legend=True):
+    """
+    
+    """
+    for result, label, color in zip(results, labels, colors):
+        if result is None:
+            continue
+        metric = get_rel_residual(result)
+        iters = (metric > rtol).sum(dim=0)
+        ax.hist(iters, edgecolor=color, facecolor=color, bins=bins, label=label)
+    if log_scale:
+        ax.set_xlim(max(xmin, 1), xmax)
+        ax.set_xscale("log")
+    else:
+        ax.set_xlim(xmin, xmax)
+    ax.set_axisbelow(True)
+    ax.grid()
+    if legend:
+        ax.legend()
+
+def get_rel_residual(result):
+    """
+    
+    """
+    return result['err_history'].cpu().flatten(start_dim=-2) / result['err_history'].cpu().flatten(start_dim=-2)[0]
+
+
+def plot_deformed_rve_2d(
+    problem,
+    disp,
+    field,
+    loading=None,
+    fluctuation_scaling: float = 1.0,
+    deformation_scaling: float = 1.0,
+    plot_loading: bool = False,
+    plot_boundary: bool = False,
+    shading: str = "gouraud",
+    file: Optional[str] = None,
+    vmin: Optional[List[float]] = None,
+    vmax: Optional[List[float]] = None,
+    figsize: Optional[List[float]] = None
+):
+    """
+
+    :param disp:
+    :param field:
+    :param loading:
+    :param fluctuation_scaling:
+    :param deformation_scaling:
+    :param plot_loading:
+    :param plot_boundary:
+    :param shading:
+    :param file:
+    :param vmin:
+    :param vmax:
+    :return:
+    """
+    def apply_mask(triang, alpha=0.4):
+        # Mask triangles with sidelength greater than a threshold alpha
+        triangles = triang.triangles
+        # Mask off unwanted triangles.
+        x = triang.x
+        y = triang.y
+        xtri = x[triangles] - np.roll(x[triangles], 1, axis=1)
+        ytri = y[triangles] - np.roll(y[triangles], 1, axis=1)
+        maxi = np.max(np.sqrt(xtri**2 + ytri**2), axis=1)
+        # apply masking
+        triang.set_mask(maxi > alpha)
+    
+    if loading is None:
+        loading = torch.eye(3, dtype=disp.dtype, device=disp.device)
+
+    n_loadings = loading.shape[0]
+
+    #if problem.n_dim == 2:
+    #    disp = disp.transpose(-1, -2)
+
+    deformed_coords = problem.get_deformed_coordinates(
+        disp, loading, fluctuation_scaling=fluctuation_scaling, deformation_scaling=deformation_scaling
+    )
+    loading_coords = problem.get_deformed_coordinates(
+        disp, loading, fluctuation_scaling=0.0, deformation_scaling=deformation_scaling
+    )
+    coords = problem.get_node_coords()
+    #boundary_idx = problem.get_boundary_idx()
+    #boundary = coords[..., boundary_idx[0], boundary_idx[1]]
+    #deformed_boundary = deformed_coords[..., boundary_idx[0], boundary_idx[1]]
+    #loading_boundary = loading_coords[..., boundary_idx[0], boundary_idx[1]]
+
+    loadings = [
+        rf"$\bar{{\boldsymbol{{\varepsilon}}}}=\bar{{\boldsymbol{{\varepsilon}}}}^{{({i + 1})}}$"
+        for i in range(n_loadings)
+    ]
+
+    if figsize is None:
+        figsize = [6.3, 2.0]
+
+    if problem.n_dim == 3:
+        N_cut = problem.n_grid[-1] // 2
+        field = field[..., N_cut]
+        deformed_coords = deformed_coords[..., N_cut]
+            
+    fig, ax = plt.subplots(1, n_loadings, figsize=figsize, dpi=300, squeeze=False)
+    for load_idx, load_name in enumerate(loadings):
+        ax[0, load_idx].axis("off")
+        ax[0, load_idx].set_aspect("equal")
+        ax[0, load_idx].set_title(load_name)
+        tri = Triangulation(deformed_coords[load_idx, 0].ravel(), deformed_coords[load_idx, 1].ravel())
+        apply_mask(tri, alpha=0.02)
+
+        if vmin is None:
+            vmin_idx = field[load_idx].min().item()
+        else:
+            vmin_idx = vmin[load_idx]
+        if vmax is None:
+            vmax_idx = 0.5 * field[load_idx].max().item()
+        else:
+            vmax_idx = vmax[load_idx]
+        
+        tpc = ax[0, load_idx].tripcolor(
+            tri, field[load_idx].ravel(), cmap="jet", shading=shading, rasterized=True, vmin=vmin_idx, vmax=vmax_idx
+        )
+
+        if plot_loading:
+            ax[0, load_idx].plot(
+                loading_boundary[load_idx, 0].ravel(),
+                loading_boundary[load_idx, 1].ravel(),
+                "k--",
+                lw=1,
+            )
+        if plot_boundary:
+            ax[0, load_idx].plot(
+                deformed_boundary[load_idx, 0].ravel(),
+                deformed_boundary[load_idx, 1].ravel(),
+                "r--",
+                lw=1,
+            )
+        clb = fig.colorbar(tpc)
+        clb.ax.set_title(r"$||\boldsymbol{\sigma}|| \,[\mathrm{GPa}]$")
+    plt.tight_layout()
+    if file is not None:
+        plt.savefig(file, dpi=300)
+    plt.show()
+
+
+def plot_deformed_rve_3d(
+    problem,
+    disp,
+    field,
+    loadings,
+    fluctuation_scaling: float = 1.0,
+    deformation_scaling: float = 1.0,
+    file: Optional[str] = None,
+    vmin: Optional[List[float]] = None,
+    vmax: Optional[List[float]] = None,
+    figsize: Optional[List[float]] = None
+):
+    disp = torch.nn.functional.pad(disp, pad=[0,1,0,1,0,1], mode="circular")
+    coords = problem.get_node_coords().to(dtype=disp.dtype, device=disp.device)
+    deformations = problem.get_deformations(disp.transpose(-1, -3), loadings, fluctuation_scaling=1.0)[0]
+    
+    x, y, z = coords[2].cpu().numpy(), coords[1].cpu().numpy(), coords[0].cpu().numpy()
+    grid = pv.StructuredGrid(x, y, z)
+    grid['vectors'] = deformations.flatten(start_dim=1).T.cpu().numpy() * deformation_scaling
+    warped = grid.warp_by_vector()
+    
+    pl = pv.Plotter()
+    pl.add_mesh(warped, scalars=field.numpy().ravel(), clim=[vmin, vmax], label="stress norm", cmap="jet", lighting=True, diffuse=0.2, specular=1.0, ambient=0.6, scalar_bar_args={"vertical": True})
+    pl.screenshot(filename=file, window_size=figsize)
+    pl.show()

{unocg-0.0.2 → unocg-0.0.4/unocg.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unocg
-Version: 0.0.2
+Version: 0.0.4
 Summary: UNO-CG
 Author-email: Julius Herb <herb@mib.uni-stuttgart.de>, Felix Fritzen <fritzen@mib.uni-stuttgart.de>
 License-Expression: MIT
@@ -20,9 +20,9 @@ Requires-Dist: torch==2.10.0
 Requires-Dist: tqdm~=4.67.1
 Requires-Dist: torchvision~=0.25.0
 Requires-Dist: pyvista[jupyter]~=0.46.5
+Requires-Dist: jupyterlab~=4.5.3
 Provides-Extra: dev
 Requires-Dist: jupytext~=1.16.4; extra == "dev"
-Requires-Dist: jupyterlab~=4.5.3; extra == "dev"
 Requires-Dist: pre_commit~=3.8.0; extra == "dev"
 Requires-Dist: sphinx~=7.4.7; extra == "dev"
 Requires-Dist: sphinx-rtd-theme~=2.0.0; extra == "dev"

{unocg-0.0.2 → unocg-0.0.4}/unocg.egg-info/SOURCES.txt

@@ -28,6 +28,15 @@ unocg/problems/thermal.py
 unocg/solvers/__init__.py
 unocg/solvers/base.py
 unocg/solvers/torch.py
+unocg/training/__init__.py
+unocg/training/losses/__init__.py
+unocg/training/losses/base.py
+unocg/training/losses/mechanical.py
+unocg/training/losses/thermal.py
 unocg/transforms/__init__.py
 unocg/transforms/base.py
-unocg/transforms/fourier.py
+unocg/transforms/fourier.py
+unocg/utils/__init__.py
+unocg/utils/data.py
+unocg/utils/evaluation.py
+unocg/utils/plotting.py

{unocg-0.0.2 → unocg-0.0.4}/unocg.egg-info/requires.txt

@@ -6,13 +6,13 @@ torch==2.10.0
 tqdm~=4.67.1
 torchvision~=0.25.0
 pyvista[jupyter]~=0.46.5
+jupyterlab~=4.5.3
 
 [all]
 unocg[dev]
 
 [dev]
 jupytext~=1.16.4
-jupyterlab~=4.5.3
 pre_commit~=3.8.0
 sphinx~=7.4.7
 sphinx-rtd-theme~=2.0.0