ucl-bug / helmnet Goto Github PK

Explore exporting the pre-trained model using the Open Neural Network Exchange (ONNX) format. This could then be imported into, e.g., MATLAB so the model can be directly run. One possible complication is the use of the FFT within the residual calculation, which may not be supported.

When I want to change the location of the source, for example from (82,48) to (48,48), I just changed the parameters in base.json. Then I ran the evaluate.py. However, from the prediction results, I found that the location of the source did not change, still (82,48). I am wondering what else I need to do when I want to change the location of the source, besides changing the parameters in base.json, before I run evaluate.py? What should I do If I want to test the situation of two sources or more?

Thanks a lot for your suggestions.

Currently, running inference fails when setting solver.to("cuda:0") because the sigma values in IterativeSolver.set_laplacian are not transferred to the GPU. The fix is to add .to(self.device) to sigmax and sigmay. The code is already there, but commented out from a previous PR.

Need a test that checks that the weights can load after training (see #21 )

The network was initially trained assuming constant density and no absorption. Initial tests show that the network may also generalise (even without re-training) to more general cases. This requires updating IterativeSolver.get_residual.

Absorption could possibly be implemented by adding an imaginary part to k_sq, while density gradients could be implemented using, e.g.,

rho_dx = self.Lap.dx(rho)

The additional parameters will need to be propagated through the code, e.g., by adding additional channels to the speed of sound input. This will need changing IterativeSolver.get_initials to only use the sound speed channel in the calculation of k_sq.

  k_sq, wavefield = self.get_initials(sos_maps)
  self.f.clear_states(wavefield)
  residual = self.get_residual(wavefield, k_sq)

To:

k_sq, wavefield, alpha, rho = self.get_initials(medium_maps)
self.f.clear_states(wavefield)
residual = self.get_residual(wavefield, k_sq, alpha, rho)

This will need to propagated in several places that call get_initials , get_residual, or use sos.

• I'm submitting a ...

  • bug report
  • feature request
  • other (e.g. discussion of ideas)

• What is the current behavior?
  We don't have any test at all. At least having a test with CI that checks that the example in the README can be run would be great, since we are slowly changing this repo

• What is the expected behavior?
  CI that tests at every merge on the main

A 3D version of the helmnet should be considered. This would require the following:

• Extending the unet to 3D (Conv3D, MaxPool3D, etc)
• Extend the training data to 3D
• Add extra dimension (z) to all inputs, PML components etc
• Extend operators (e.g., Spectral.LaplacianWithPML, etc)
• Potentially reduce batch size for training
• Check the replay buffer works with 3D data (or could potentially be replaced with binomial checkpointing as now being implemented in j-Wave)

Feature request

• What is the current behavior?
  Currently, there's no way to set the source location when running evaluate.py (see #9 ).

• What is the expected behavior?
  To be able to set the source location while running the evaluation script

• Other information
  The get_model method accepts source_location it as input. Adding the functionality should be as easy as adding a new argument to argparse and pass it to the __init__ method, down to get_model.

For compatibility, the version of pytorch-lightning in environment.yml should be fixed, e.g., to 1.2.10.

According to this repository for Fourier Neural Operators, they've made their code 30% faster by switching to native complex representation in PyTorch, which wasn't available when this work was published.

Is it worth refactoring the code to use complex multiplications? On one hand, it would be nice to have a faster codebase, on the other hand this repo makes sense for research but not really for production, so perhaps one doesn't really care about speed.

@sio13 @SonyPony Any suggestion? (Is it about time to just move everything to JAX? 😄 )