Merge branch 'main' into devel

EXPtools/utils/coefs_index.py

@@ -0,0 +1,80 @@
+"""
+This module provides functions to work with index-based calculations for spherical harmonics coef series.
+"""
+import numpy as np, math
+
+def total_terms(l_max):
+    """
+    Calculate the total number of terms up to a given maximum angular number.
+
+    Parameters:
+        l_max (int): The maximum angular quantum number.
+
+    Returns:
+        float: The total number of terms.
+    """
+    return l_max * (l_max + 1) / 2 + l_max + 1  # Sigma(lmax+1) + 1 (the extra one for the 0th term)
+
+def I(l, m):
+    """
+    Calculate the index of a spherical harmonic element given the angular numbers l and m .
+
+    Parameters:
+        l (int): The angular number.
+        m (int): The magnetic quantum number, ranging from 0 to l.
+
+    Returns:
+        int: The index corresponding to the specified angular numbers.
+    """
+    import math
+    assert isinstance(l, int) and isinstance(m, int), "l and m must be integers"
+    assert l >= 0, "l must be greater than 0"
+    assert abs(m) <= l, "m must be less than or equal to l"
+    return int(l * (l + 1) / 2) + abs(m)
+
+def inverse_I(I):
+    """
+    Calculate the angular numbers l and m given the index of a spherical harmonic element.
+
+    Parameters:
+        I (int): The index of the spherical harmonic element.
+
+    Returns:
+        tuple: A tuple containing the angular numbers (l, m).
+    """
+    import math
+    assert isinstance(I, int) and I >=0, "I must be an interger greater than or equal to 0"
+    l = math.floor((-1 + math.sqrt(1 + 8 * I)) / 2)  # Calculate l using the inverse of the formula
+    m = I - int(l * (l + 1) / 2)  # Calculate m using the given formula
+    return l, m
+
+def set_specific_lm_non_zero(data, lm_pairs_to_set):
+    """
+    Sets specific (l, m) pairs in the input data to non-zero values.
+
+    Parameters:
+        data (np.ndarray): Input data, an array of complex numbers.
+        lm_pairs_to_set (list): List of tuples representing (l, m) pairs to set to non-zero.
+
+    Returns:
+        np.ndarray: An array with selected (l, m) pairs set to non-zero values.
+
+    Raises:
+        ValueError: If any of the provided (l, m) pairs are out of bounds for the input data.
+    """
+    assert isinstance(lm_pairs_to_set, list), "lm_pairs_to_set must be a list"
+    for pair in lm_pairs_to_set:
+        assert isinstance(pair, tuple), "Each element in lm_pairs_to_set must be a tuple"
+        assert len(pair) == 2, "Each tuple in lm_pairs_to_set must contain two elements"
+        assert all(isinstance(x, int) for x in pair), "Each element in each tuple must be an integer"
+
+    # Get a zeros array of the same shape as the input data
+    arr_filt = np.zeros(data.shape, dtype=complex)
+
+    # Check if the provided (l, m) pairs are within the valid range
+    for l, m in lm_pairs_to_set:
+        if l < 0 or l >= data.shape[0] or m < 0 or m > l:
+            raise ValueError(f"Invalid (l, m) pair: ({l}, {m}). Out of bounds for the input data shape.")
+        arr_filt[I(l, m), :, :] = data[I(l, m), :, :] 
+
+    return arr_filt

EXPtools/visuals/visualize.py

@@ -139,8 +139,8 @@ def spherical_avg_prop(basis, coefficients, time=0, radius=np.linspace(0.1, 600,
     return np.array(field), radius
 
 def return_fields_in_grid(basis, coefficients, times=[0], 
-                          projection='3D', proj_plane=0, 
-                          grid_lim=300, npoints=150,):
+                       projection='3D', proj_plane=0, 
+                       grid_lim=300, npoints=150,):
 
     """
     Returns the 2D or 3D field grids as a dict at each time step as a key.
@@ -150,11 +150,11 @@ def return_fields_in_grid(basis, coefficients, times=[0],
     coefficients (obj): object containing the coefficients for the simulation
     times (list): list of the time to compute the field grid.
     projection (str): the slice projection to plot. Can be '3D', XY', 'XZ', or 'YZ'.
-    proj_plane (float, optional): the value of the coordinate that is held constant in the 2D slice projection
-    npoints (int, optional): the number of grid points in each dimension
+    proj_plane (float, optional): the value of the coordinate that is held constant in the 2D slice projection.
+    npoints (int, optional): the number of grid points in each dimension. It's +1 in 3D projection.
     grid_limits (float, optional): the limits of the grid in the dimensions.
 
-    Returns: fields along with the specified 2D or 3D grid. The fields are in a
+    Returns:
     heirarchical dict structure with for each time:
         dict_keys(['d', 'd0', 'd1', 'dd', 'fp', 'fr', 'ft', 'p', 'p0', 'p1'])
         where 
@@ -168,33 +168,37 @@ def return_fields_in_grid(basis, coefficients, times=[0],
         - 'p': Total potential at each point.
         - 'p0': Potential in the monopole term.
         - 'p1': Potential in l>0 terms.
+        
+    xgrid used to make the fields. 
     """
 
     assert projection in ['3D', 'XY', 'XZ', 'YZ'], "Invalid value for 'projection'. Must be one of '3D', 'XY', 'XZ', or 'YZ'."
-
-
 
     if projection == '3D':
         pmin  = [-grid_lim, -grid_lim, -grid_lim]
         pmax  = [grid_lim, grid_lim, grid_lim]
         grid  = [npoints, npoints, npoints]
 
         ##create a projection grid along with it
-        x = np.linspace(-grid_lim, grid_lim, npoints)
+        x = np.linspace(-grid_lim, grid_lim, npoints+1) ##pyEXP creates a grid of npoints+1 for 3D spacing.
         xgrid = np.meshgrid(x, x, x)
+
+        field_gen = pyEXP.field.FieldGenerator(times, pmin, pmax, grid)
+        return field_gen.volumes(basis, coefficients), xgrid
 
     else:
         x = np.linspace(-grid_lim, grid_lim, npoints)
         xgrid = np.meshgrid(x, x)
-        grid  = [npoints, npoints]
-
+
         if projection == 'XY':        
             pmin  = [-grid_lim, -grid_lim, proj_plane]
             pmax  = [grid_lim, grid_lim, proj_plane]
+            grid  = [npoints, npoints, 0]
 
         elif projection == 'XZ':
             pmin  = [-grid_lim,  proj_plane, -grid_lim]
             pmax  = [grid_lim, proj_plane, grid_lim]
+            grid  = [npoints, 0, npoints]
 
         elif projection == 'YZ':
             pmin  = [proj_plane,  -grid_lim, -grid_lim]
@@ -325,7 +329,4 @@ def slice_3d_fields(basis, coefficients, time=0,  npoints=50,
 
     if prop == 'force':
         return [fx.reshape(npoints, npoints, npoints), fy.reshape(npoints, npoints, npoints), fz.reshape(npoints, npoints, npoints)], xgrid
-
-
-
-
+

README.md

@@ -11,4 +11,4 @@ python -m pip install .
 
 ### Docuentation:
 
-The ``EXPtools`` documenntation is hosted (temporaily) [here](https://users.flatironinstitute.org/~nico/EXPtools/docs/) 
+The ``EXPtools`` documenntation is hosted (temporarily) [here](https://users.flatironinstitute.org/~nico/EXPtools/docs/)