ucl-bug / jwave

A JAX-based research framework for differentiable and parallelizable acoustic simulations, on CPU, GPUs and TPUs
GNU Lesser General Public License v3.0
143 stars 21 forks source link

Helmholtz Problem tutorial NotFoundLookupError #215

Closed elma16 closed 1 year ago

elma16 commented 1 year ago

Describe the bug Running the helmholtz problem tutorial produces the following error NotFoundLookupError: For function "helmholtz", signature Signature(jaxdf.discretization.FourierSeries, jwave.geometry.MediumType[<class 'float'>, <class 'float'>, <class 'float'>], builtins.float) could not be resolved. which occurs at the line where params = helmholtz.default_params(src, medium, omega) is defined.

To Reproduce Steps to reproduce the behavior:

  1. Running the first five cells of the notebook produces the error.

Expected behavior A clear and concise description of what you expected to happen. I was expecting j-Wave to solve the helmholtz equation :)

Desktop (please complete the following information):

Additional context Add any other context about the problem here. N/A

astanziola commented 1 year ago

Sorry for this @elma16 . If it is referring to the helmholtz_solver_differentiable notebook, it should be fixed with the release I just made. Do you mind checking if it is still an issue, after updating? If not can you tell me in which of the two notebooks is happening? Thanks!

astanziola commented 1 year ago

Ah no there's indeed an error in one of the notebook, the helmholtz operator does not dispatch on omega, as it is always supposed to be a float number (I don't think it makes sense to have an Helmholtz operator where omega is something else).

So the way it must be called is helmholtz(src, medium, omega=omega). I'm fixing it now and will push the updated notebook

astanziola commented 1 year ago

The new updated notebook should not have this problem anymore. Feel free to reopen if that's not the case!

elma16 commented 1 year ago

Hi @astanziola - yes, the fix helmholtz(src, medium, omega=omega) indeed solves the problem I was describing. Thanks! I was also wondering why there is a difference between the presentation of the helmholtz equation in this notebook (with a source term $S_M$) and the second-order wave equation (in section 2.2, page 8) after fourier transforming the solution - which has source terms $S_M$ and $S_F$? Is the third and fourth term within the bracket somehow related to $S_F$?

astanziola commented 1 year ago

I guess you are talking about the k-Wave manual? In that case, the jwave one simply only has the mass source term 😁

In theory, all you need is a single scalar complex field, which is normally a pressure field, to define the r.h.s. the Helmholtz equation. Note that in the second order wave equation the force field is casted to a scalar field by the divergence operator.

The reason why the r.h.s. of the helmholtz operator expects a mass source in jwave is just a convention. You can use a pressure field as source just by dividing it by $i\omega$, and you can use a force source just by doing $\rho_0 \nabla \cdot S_F/i\omega$.

elma16 commented 1 year ago

Oh, sorry - I forgot to read the message properly and didn't mention that I was talking about the k-wave manual. Good reading between the lines there! In which case, does my reasoning make sense here?

Starting from the second order wave equation with source terms as in the manual:

$$c^2 \nabla^2 p - \frac{\partial^2 p}{\partial t^2} = \rho_0 c^2 \nabla \cdot S_F - \frac{\partial}{\partial t} S_M$$

Converting this to the corresponding helmholtz equation we have

$$c^2 \nabla^2 P + \omega^2 P = \hat{\rho}_0 c^2 \nabla \cdot \hat{S}_F - c^2 i \omega \hat{S}_M$$

In the notebook we have

$$(\omega^2 + c^2 \nabla^2 - \frac{c^2}{\rho_0} \nabla \rho_0 \cdot \nabla + 2i\omega^3 \alpha_0 c)P = i\omega c^2 S_M$$

So,

$$c^2 \nabla^2 P + \omega^2 P = (\frac{c^2}{\rho_0} \nabla \rho_0 \cdot \nabla) P - 2i\omega^3 \alpha_0 c P + i \omega c^2 S_M$$

From this, can we deduce that the first two terms on the RHS of the "notebook helmholtz" are the same as the first term on the RHS of the "manual helmholtz"?

astanziola commented 1 year ago

Mmm not quite,I am not sure why you have the $c^2$ term only in front of the force term and not also in front of the mass term.

But in any case, the reasoning we followed is like this (I am ignoring attenuation because that is irrelevant for what we are saying here, and it matches what you wrote)

Starting from (2.8) of the k-wave manual, if you Fourier-transform in time you get the system of equations

$$ i\omega \mathbf{u} = \frac{1}{\rho_0}\nabla p + \mathbf{S}_F $$

$$ \frac{i\omega}{c^2}p = \rho_0\nabla \cdot \mathbf{u} + S_M $$

Plugging the first into the second one gives

$$ \frac{i\omega}{c^2}p = \rho_0 \nabla \cdot \left[ \frac{1}{i\omega} \left( \frac{1}{\rho_0}\nabla p + \mathbf{S}_F \right) \right] + S_M $$

which simplifies to

$$ \frac{\omega^2}{c^2}p = - \rho_0 \nabla \cdot \left( \frac{1}{\rho_0} \nabla p \right) - \rho_0 \nabla \cdot \mathbf{S}_F + i\omega S_M $$

Now, using $\nabla \cdot (f \nabla g) = f\nabla^2 g + \nabla f \cdot \nabla g$ and $f\nabla \frac{1}{f} + \frac{1}{f}\nabla f = 0$ we can rewrite it as

$$ \nabla^2 p + \frac{\omega^2}{c^2}p - \frac{1}{\rho_0}\nabla \rho_0 \cdot \nabla p = - \rho_0 \nabla \cdot \mathbf{S}_F + i\omega S_M $$

If we add the attenuation term we get

$$ \left(\nabla^2 + \frac{\omega^2}{c^2} - \frac{1}{\rho_0}\nabla \rho_0 \cdot \nabla + \frac{2i\omega^3\alpha_0}{c_0} \right) p = - \rho_0 \nabla \cdot \mathbf{S}_F + i\omega S_M $$

So as you see, the mass term is simply not included in jwave. This is however without any loss of generaly, since I can always set

$$ S_M = -\frac{\rho_0}{i\omega} \nabla \cdot S_F $$

to get an equivalent mass term for any source term.

If you want to look it in another way, you can think about it like this

$$ S_M^{jwave} = S_M^{kwave} -\frac{\rho_0}{i\omega} \nabla \cdot S_F $$

Does this make any sense?

astanziola commented 1 year ago

Oh, one more thing that may be creating confusion. I believe that 2.9 in the k-wave manual is assuming homogeneous density.

elma16 commented 1 year ago

I see! This makes a lot more sense now - yes I definitely made a typo in my above comment. Thank you for clearing this all up for me!

astanziola commented 1 year ago

No worries at all, let me know if there's anything else that sparks doubts :)

elma16 commented 1 year ago

Looking at your calculation again, I'm wondering if you made a typo with the governing equations, since there should be a minus sign on the first term of the RHS of both of these equations. Of course, everything works out ok, but I get a $-i \omega S_M$ - is this correct?

astanziola commented 1 year ago

Oh yes sorry! When I copied them from the notebook I simplified the minus signs in 2.8 but forgot to put a minus in front of the source terms