As a requirement of JAX and the optimizer, the loss function needs to be continuously differentiable. In our case, we want to target the membrane lateral density profile, which is typically obtained by computing the histogram of the z coordinate of the particles (assuming here that z is the direction normal to the membrane). To calculate a differentiable density profile, we can approximate the histogram with a simple Gaussian kernel density estimation, where we center a normal distribution around each discrete value z i :