wheelermt / pygpe

Library for numerically solving the Gross-Pitaevskii equations for scalar, two-component, spin-1, and spin-2 Bose-Einstein condensate systems.

Hey, I'm Matt Wheeler

---

Welcome to my GitHub!

I am a Senior Modelling & Simulation Engineer, working mainly with Fortran, C, and Matlab. In my spare time I love all-things systems programming, where I devote a lot of my time to learning and using Rust.

---

Personal Stuff:

• 👨🏻‍💻 I’ve created BEC++ and PyGPE: fast, adaptable Gross-Pitaevskii equation solvers built using C++, Python & CUDA.
• 🎐 I am also working on a quantum vortex tracker built using deep learning!
• 🌱 I’m currently learning about the Rust programming language.
• 💬 Ask me about my research!
• 📫 How to reach me: [email protected].
• 😄 Pronouns: he/him.

📈 My GitHub contributions for the past year:

📊 How I've been spending my last 30 days:

From: 12 July 2024 - To: 11 August 2024

Total Time: 5 hrs 35 mins

Markdown         1 hr 37 mins    >>>>>>>------------------   29.16 %
C#               1 hr 13 mins    >>>>>--------------------   21.87 %
Other            1 hr 10 mins    >>>>>--------------------   20.95 %

pygpe's Issues

Vortices module failing to create phase profile

There is a problem in which the vortex_phase_profile function within the vortices module is mixing arrays defined on both the CPU and GPU.
For example, calling vortex_phase_profile function(grid, 100, 1.0) results in the traceback

Attempting to find 100 positions...
Successfully found 100 positions in 100 iterations!
Traceback (most recent call last):
  File "c:\dev\python\pygpe\examples\scalar\2d_SQV_imprinting.py", line 18, in <module>
    phase = vortex_phase_profile(grid, 100, 1.0)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\dev\python\pygpe\pygpe\shared\vortices.py", line 100, in vortex_phase_profile
    heaviside_x_plus = _heaviside(x_plus)
                       ^^^^^^^^^^^^^^^^^^
  File "C:\dev\python\pygpe\pygpe\shared\vortices.py", line 61, in _heaviside
    return np.where(array < 0, np.zeros(array.shape), np.ones(array.shape))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<__array_function__ internals>", line 200, in where
  File "cupy\_core\core.pyx", line 1723, in cupy._core.core._ndarray_base.__array_function__
  File "C:\dev\python\pygpe\.venv\Lib\site-packages\cupy\_sorting\search.py", line 211, in where
    return _where_ufunc(condition.astype('?'), x, y)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "cupy\_core\_kernel.pyx", line 1287, in cupy._core._kernel.ufunc.__call__
  File "cupy\_core\_kernel.pyx", line 160, in cupy._core._kernel._preprocess_args
  File "cupy\_core\_kernel.pyx", line 146, in cupy._core._kernel._preprocess_arg
TypeError: Unsupported type <class 'numpy.ndarray'>

The error lies within the module using NumPy functions mixed with the GPU arrays of the Grid class. I propose to just swap all NumPy functions within the vortices module over to CuPy versions.

Update PyGPE dependencies to use NumPy 2.0

Feature Request

Is your feature request related to a problem? Please describe.
NumPy 2.0 is out, bringing with it a whole host of new features and improvements. From first glance, none of the breaking changes should affect PyGPE, and therefore it should be a painless upgrade.

Describe the solution you'd like
Add the latest NumPy version to pyproject.toml using poetry: poetry add numpy@latest.

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Teachability and Documentation
README.md should be updated to reflect the new dependency.

Update examples to generate animated plots

Feature Request

Rather than static plots showing one frame of data, why not make use of all the frames of data generated in the examples and plot a movie of the evolution. This showcases the capabilities of the library better, and is overall more interesting.

Solution

The easiest way to do this is using Matplotlib's [FuncAnimation}(https://matplotlib.org/stable/api/_as_gen/matplotlib.animation.FuncAnimation.html) class, meaning we can rely on dependencies already specified for the library. To do the animation, we can create an imshow of, e.g., the atom density at each frame, and stitch the result into a .gif file. The initial size of the grids may have to be lowered to, e.g., 256 x 256 if the animation takes too long to render.

Refactor DataManager class into single class

Currently, there is a lot of repeated code for the DataManager class in each system. Instead, I propose to develop a single DataManager class that handles the correct saving of data depending on the type of the wave function passed in. The api will still function in the same way. Additionally, this will remove a lot of redundant information in the documentation.

Automatically apply FFT after any initial state wavefunction changes

Feature Request

Currently a user has to manually call wfn.fft() to populate the k-space wavefunction arrays after changing the initial state. I don't like this idea as it can lead to simulations which produce garbage results. Instead, I believe that if the user modifies the (position-space) wavefunction arrays through the use of the provided methods (e.g., apply_phase or set_wavefunction) the k-space arrays should update automatically by internally calling fft().

Solution

For every function that modifies the position-space arrays of the Wavefunction object, call self.fft() at the end of each respective method.

Alternatives

Should it be possible to bypass this behaviour? This could be achieved by having a default boolean to each function call like: overwrite_kspace=True which the user can override to False if for whatever reason the k-space arrays should not be updated.

Documentation updates

The function signature of each function should be updated if the overwrite_kspace=False method is used. Additionally, all examples in the examples/ folder should be updated to not explicitly call wfn.fft() after setting the initial state.

Remove redundant `asarray` calls

There are redundant cp.asarray calls in grid.py when defining the wavenumber array:

pygpe/pygpe/shared/grid.py

Line 109 in 1474e17

self.wave_number = cp.asarray(self.fourier_x_mesh**2)

and in both examples in the examples/ folder:

pygpe/examples/scalar/2d_SQV_imprinting.py

Line 25 in 1474e17

psi.apply_phase(cp.asarray(phase)) # Apply phase to wavefunction

These are remnants from an earlier version of PyGPE, and are no longer needed.

Implement loading capabilities of `DataManager`

Feature Request

Currently, the DataManager classes can only save data, they cannot load it. When working on #19 I found I was having to manually load the data saved by a DataManager using H5py. Instead, a DataManager should have a load_data method which simply returns a h5py.File instance of the data file, which can then be accessed.

Since this will be common to all DataManager instances it should be implemented in the abstract _DataManager class.

Considerations

This should have appropriate error checks which catch if the user has not already specified the filename and data_path attributes.

Refactor `Wavefunction` classes to use inheritance

Currently, the Wavefunction class for each system are completely separate entities. It would make more sense to have a base Wavefunction abstract class which specifies the shared methods between all subclasses of Wavefunction, which then is inherited by the Wavefunction class of each system. For example, the Wavefunction class for the spin-1 system would be renamed, e.g., SpinOneWavefunction and implement the correct way of adding noise, applying phases etc.

This in turn makes the code more readable, and should help in refactoring other parts of the codebase down the line.

Meshgrid order issue when pre-calculating wavenumber with anisotropic grid.

Hi, when using the package I found an issue when you create grid.wavenumber.
The wavenumber is created with numpy/cupy meshgrid, but you should specify indexing='ij' instead of indexing='xy' (default) in the following code:
( self.fourier_x_mesh, self.fourier_y_mesh, self.fourier_z_mesh, ) = cp.meshgrid(fourier_x, fourier_y, fourier_z,indexing='ij')