#### NOTE New to DeepInverse? Get started with the basics with the [5 minute quickstart tutorial](https://deepinv.org/auto_examples/basics/demo_quickstart.html.md#sphx-glr-auto-examples-basics-demo-quickstart-py).. # Bring your own physics This examples shows you how to use DeepInverse with your own physics. While DeepInverse offers a [large number of forward operators](https://deepinv.org/user_guide/physics/physics.html.md#physics), you can also bring your own forward operator for your specific imaging problem. DeepInverse’s modular [physics framework](https://deepinv.org/user_guide/physics/intro.html.md#physics-intro) makes this easy by letting you inherit useful methods from the most appropriate physics base class: * [`deepinv.physics.Physics`](https://deepinv.org/api/stubs/deepinv.physics.Physics.html.md#deepinv.physics.Physics) for non-linear operators; * [`deepinv.physics.LinearPhysics`](https://deepinv.org/api/stubs/deepinv.physics.LinearPhysics.html.md#deepinv.physics.LinearPhysics) for linear operators; * [`deepinv.physics.DecomposablePhysics`](https://deepinv.org/api/stubs/deepinv.physics.DecomposablePhysics.html.md#deepinv.physics.DecomposablePhysics) for linear operators with a closed-form singular value decomposition. #### SEE ALSO Often your physics can be modelled without much work. You could: * Use an existing physics but with custom `params`, e.g. blur with a custom kernel, or MRI with a custom sampling pattern. See [parameter dependent operators](https://deepinv.org/user_guide/physics/intro.html.md#parameter-dependent-operators). * Inherit from an existing physics but override or wrap a particular method e.g. [`deepinv.physics.MRI`](https://deepinv.org/api/stubs/deepinv.physics.MRI.html.md#deepinv.physics.MRI) $\rightarrow$ [`deepinv.physics.DynamicMRI`](https://deepinv.org/api/stubs/deepinv.physics.DynamicMRI.html.md#deepinv.physics.DynamicMRI) * Define a new operator by [combining existing operators](https://deepinv.org/user_guide/physics/intro.html.md#physics-combining). In this example we will demonstrate creating a simple forward operator from scratch that converts RGB images to grayscale images. We also show you how to exploit the singular value decomposition of the operator to speed up the evaluation of the pseudo-inverse and proximal operators. ```Python from __future__ import annotations import deepinv as dinv import torch ``` ## Creating a custom forward operator. Defining a new linear operator only requires a forward function $\forw{\cdot}$ and its adjoint operation $A^\top(\cdot)$, inheriting the remaining structure of the [`deepinv.physics.LinearPhysics`](https://deepinv.org/api/stubs/deepinv.physics.LinearPhysics.html.md#deepinv.physics.LinearPhysics) class. Once the operator is defined, we can use any of the functions in the [`deepinv.physics.Physics`](https://deepinv.org/api/stubs/deepinv.physics.Physics.html.md#deepinv.physics.Physics) class and [`deepinv.physics.LinearPhysics`](https://deepinv.org/api/stubs/deepinv.physics.LinearPhysics.html.md#deepinv.physics.LinearPhysics) class, such as computing the norm of the operator, testing the adjointness, computing the proximal operator, etc. #### TIP By default, the adjoint of a [`LinearPhysics`](https://deepinv.org/api/stubs/deepinv.physics.LinearPhysics.html.md#deepinv.physics.LinearPhysics) is computed using autograd with [`deepinv.physics.adjoint_function`](https://deepinv.org/api/stubs/deepinv.physics.adjoint_function.html.md#deepinv.physics.adjoint_function). Note however that defining a closed form adjoint is generally more computationally efficient in memory and time. #### NOTE To make the new physics compatible with other torch functionalities, all physics parameters (i.e. attributes of type [`torch.Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor)) should be registered as [module buffers](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer) by using `self.register_buffer(param_name, param_tensor)`. This ensures methods like `.to(), .cuda()` work properly, allowing one to train a model using Distributed Data Parallel. #### IMPORTANT As per [`torch.nn.Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module) functionalities, it is expected that all learnable parameters and registered buffers of the physics are on the same device and have the same data type. To make the initialization of the physics more flexible and efficient, you may want to pass a `device` and `dtype` argument to the physics constructor, and use them to initialize the physics parameters/buffers. It is important that these arguments are not set as instance variables, e.g. `self.device = device`, since `self.device` is not automatically updated when calling `physics.to(device)`. #### NOTE Below, when calling `super().update_parameters(coefficients=coefficients, **kwargs)` in `Decolorize.A`, the `super().update_parameters` method takes care of pushing the incoming parameters to the correct device and data type. For efficiency, it is nonetheless recommended to initialize the incoming `coefficients` parameters on the correct device, to avoid unnecessary data transfers. #### TIP Inherit from [mixin](https://deepinv.org/user_guide/other/utils.html.md#mixin) classes to provide specialized methods for your physics. ```Python class Decolorize(dinv.physics.LinearPhysics): r""" Converts RGB images to grayscale. """ def __init__( self, device: str | torch.device = "cpu", dtype: torch.dtype = torch.float32, **kwargs, ): super().__init__(**kwargs) coefficients = torch.tensor( [0.2989, 0.5870, 0.1140], dtype=dtype, device=device ) self.register_buffer("coefficients", coefficients) def A( self, x: torch.Tensor, coefficients: torch.Tensor = None, **kwargs ) -> torch.Tensor: """Forward operator. :param torch.Tensor x: input image with 3 color (RGB) channels, i.e. [*,3,*,*] :param torch.Tensor coefficients: optionally set coefficients on the fly :param dict kwargs: any other keyword parameters to set on the fly, such as noise model sigma :return: torch.Tensor grayscale measurements """ super().update_parameters(coefficients=coefficients, **kwargs) y = x * self.coefficients[None, :, None, None] return torch.sum(y, dim=1, keepdim=True) ``` ## Test physics We test our physics on a toy image with 3 color channels. We add a Gaussian noise model to the linear physics. We simulate measurements using the forward operator $y=\forw{x}+\epsilon$. We then leverage the linear physics base class to automatically compute the linear pseudo-inverse $A^\dagger y$. ```Python x = torch.zeros(1, 3, 96, 128) x[:, 0, :32, :] = 1 x[:, 1, 32:64, :] = 1 x[:, 2, 64:, :] = 1 physics = Decolorize( img_size=(3, 96, 128), noise_model=dinv.physics.GaussianNoise(sigma=0.1) ) y = physics(x) dinv.utils.plot({"x": x, "y": y, "Linear pseudo-inverse": physics.A_dagger(y)}) ``` It is often useful for reconstruction algorithms that the physics has unit norm, which you can verify using [`deepinv.physics.LinearPhysics.compute_norm()`](https://deepinv.org/api/stubs/deepinv.physics.LinearPhysics.html.md#deepinv.physics.LinearPhysics.compute_norm). We see that this physics fails this. ```Python print(f"The linear operator has norm={physics.compute_sqnorm(x):.2f}") ``` ```none Power iteration converged at iteration 2, ||A^T A||_2=0.45 The linear operator has norm=0.45 ``` All parameters or buffers of the physics, such as `coefficients` in the case of `Decolorize`, can be updated on the fly with `physics.update(**params)` or in the forward pass `physics(x, **params)`: ```Python print( "Original coefficients and sigma:", physics.coefficients, physics.noise_model.sigma ) physics.update(coefficients=torch.tensor([1.0, 2.0, 3.0]), sigma=0.2) print( "Updated coefficients and sigma via update:", physics.coefficients, physics.noise_model.sigma, ) y = physics(x, coefficients=torch.tensor([4.0, 5.0, 6.0]), sigma=0.3) print( "Updated coefficients and sigma via forward pass:", physics.coefficients, physics.noise_model.sigma, ) ``` ```none Original coefficients and sigma: tensor([0.2989, 0.5870, 0.1140]) tensor(0.1000) Updated coefficients and sigma via update: tensor([1., 2., 3.]) tensor(0.2000) Updated coefficients and sigma via forward pass: tensor([4., 5., 6.]) tensor(0.3000) ``` ## Implementing a closed form adjoint Instead, if we know the closed form of the adjoint operator, we can implement it directly in [`deepinv.physics.LinearPhysics.A_adjoint()`](https://deepinv.org/api/stubs/deepinv.physics.LinearPhysics.html.md#deepinv.physics.LinearPhysics.A_adjoint), instead of relying on [`autodifferentiation`](https://deepinv.org/api/stubs/deepinv.physics.adjoint_function.html.md#deepinv.physics.adjoint_function) which is generally less efficient. An additional benefit of implementing the adjoint is that we no longer require the `img_size` parameter when creating the operator, which was needed for autodifferentiation. If you implement your own adjoint, it is recommended to verify that it is well-defined using [`deepinv.physics.LinearPhysics.adjointness_test()`](https://deepinv.org/api/stubs/deepinv.physics.LinearPhysics.html.md#deepinv.physics.LinearPhysics.adjointness_test). ```Python class Decolorize2(Decolorize): """Override previous Decolorize using a closed-form adjoint.""" def A_adjoint( self, y: torch.Tensor, coefficients: torch.Tensor = None, **kwargs ) -> torch.Tensor: """Closed-form adjoint operator. :param torch.Tensor y: input grayscale measurements :param torch.Tensor coefficients: optionally set coefficients on the fly :param dict kwargs: any other keyword parameters to set on the fly, such as noise model sigma :return: torch.Tensor adjoint reconstruction """ super().update_parameters(coefficients=coefficients, **kwargs) return y * self.coefficients[None, :, None, None] physics = Decolorize2() if physics.adjointness_test(x) < 1e-5: print("The linear operator has a well defined adjoint") ``` ```none The linear operator has a well defined adjoint ``` ## Creating a decomposable forward operator. If the forward operator has a closed form singular value decomposition (SVD), you should instead implement the operator using the [`deepinv.physics.DecomposablePhysics`](https://deepinv.org/api/stubs/deepinv.physics.DecomposablePhysics.html.md#deepinv.physics.DecomposablePhysics) base class. The operator $A$ in this example has a known closed-form SVD: $$ A = U\text{diag}(s)V^{\top} $$ where $\text{diag}(s)$ is the `mask` of singular values. #### TIP As in the case of `LinearPhysics`, `V` and `U_adjoint` are implemented using autodifferentiation by default, but you can implement them directly if you know the closed form of the operator. ```Python class DecolorizeSVD(dinv.physics.DecomposablePhysics): r""" Converts RGB images to grayscale. We use unnormalized coefficients and a singular value of 0.447. Here, `U` and `U_adjoint` are set to the identity. """ def __init__(self, **kwargs): super().__init__(mask=0.447, **kwargs) coefficients = torch.tensor([0.6687, 1.3132, 0.2550], dtype=torch.float32) self.register_buffer("coefficients", coefficients) def V_adjoint(self, x: torch.Tensor) -> torch.Tensor: y = x * self.coefficients[None, :, None, None] return torch.sum(y, dim=1, keepdim=True) def V(self, y: torch.Tensor) -> torch.Tensor: return y * self.coefficients[None, :, None, None] physics2 = DecolorizeSVD(noise_model=dinv.physics.GaussianNoise(sigma=0.1)) y2 = physics2(x) dinv.utils.plot({"x": x, "y": y2, "Linear pseudo-inverse": physics2.A_dagger(y2)}) if physics.adjointness_test(x) < 1e-5: print("The decomposable operator has a well defined transpose") print(f"The decomposable operator has norm={physics.compute_sqnorm(x):.2f}") ``` ```none The decomposable operator has a well defined transpose Power iteration converged at iteration 2, ||A^T A||_2=0.45 The decomposable operator has norm=0.45 ``` ## Benefits of using a decomposable forward operator. The main benefit of using a decomposable forward operator is that it provides closed form solutions for the proximal operator and the linear pseudo-inverse. Moreover, some algorithms, such as [`deepinv.sampling.DDRM`](https://deepinv.org/api/stubs/deepinv.sampling.DDRM.html.md#deepinv.sampling.DDRM) require the forward operator to be decomposable. ```Python import time def sync_cuda(): if torch.cuda.is_available(): torch.cuda.synchronize() sync_cuda() start = time.time() for i in range(10): xlin = physics.A_dagger(y) xprox = physics.prox_l2(x, y, 0.1) sync_cuda() end = time.time() print(f"Elapsed time for LinearPhysics: {end - start:.2f} seconds") sync_cuda() start = time.time() for i in range(10): xlin2 = physics2.A_dagger(y) xprox2 = physics2.prox_l2(x, y2, 0.1) sync_cuda() end = time.time() print(f"Elapsed time for DecomposablePhysics: {end - start:.2e} seconds") ``` ```none Elapsed time for LinearPhysics: 0.01 seconds Elapsed time for DecomposablePhysics: 8.43e-04 seconds ``` 🎉 Well done, you now know how to implement your own physics! ### What’s next? * Check out [the example on how to inference a state-of-the-art general pretrained model](https://deepinv.org/auto_examples/basics/demo_pretrained_model.html.md#sphx-glr-auto-examples-basics-demo-pretrained-model-py) with your new physics. * Check out the [example on how to fine-tune a foundation model](https://deepinv.org/auto_examples/models/demo_foundation_model.html.md#sphx-glr-auto-examples-models-demo-foundation-model-py) to your own physics. **Total running time of the script:** (0 minutes 0.160 seconds) [`Download Jupyter notebook: demo_custom_physics.ipynb`](https://deepinv.org/auto_examples/basics/_downloads/18622f4a3fd14981b3f936fc8221271d/demo_custom_physics.ipynb) [`Download Python source code: demo_custom_physics.py`](https://deepinv.org/auto_examples/basics/_downloads/ad174c57928b9ac2c9a785d6df428128/demo_custom_physics.py) [`Download zipped: demo_custom_physics.zip`](https://deepinv.org/auto_examples/basics/_downloads/57df16b7767233f84e3f23d1d30244bd/demo_custom_physics.zip) [Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)