Going through documentation of chromatix (https://chromatix.readthedocs.io/en/latest/101/) I spotted some errors when importing functions from python scripts (maybe due to changes and missing update)
This is the original import section:
from typing import Optional
from functools import partial
import jax.numpy as jnp
from jax import random
from flax import linen as nn
from chex import Array
from einops import rearrange
import matplotlib.pyplot as plt
import chromatix.functional as cx
from chromatix import Field, ScalarField
from chromatix.systems import OpticalSystem, Microscope, Optical4FSystemPSF
from chromatix.elements import PlaneWave, FFLens, ThinSample, BasicSensor, ZernikeAberrations
from chromatix.utils import siemens_star, trainable
Instead, I needed to correct trainable, seidel_aberrations, and zernike_aberrations:
from chromatix.utils import siemens_star
from chromatix.elements import trainable
from chromatix.utils.initializers import seidel_aberrations, zernike_aberrations
Going through documentation of chromatix (https://chromatix.readthedocs.io/en/latest/101/) I spotted some errors when importing functions from python scripts (maybe due to changes and missing update)
This is the original import section:
Instead, I needed to correct trainable, seidel_aberrations, and zernike_aberrations:
Just to keep in mind and don't forget! Thank you!