Merge pull request #13 from jngaravitoc/devel

Devel

EXPtools/basis_builder/basis_utils.py

@@ -1,6 +1,6 @@
 import os, sys, pickle, pyEXP
 import numpy as np
-from . import  makemodel
+from EXPtools.basis_builder import makemodel
 
 def make_config(basis_id, numr, rmin, rmax, lmax, nmax, scale, 
                 modelname='', cachename='.slgrid_sph_cache'):
@@ -242,18 +242,15 @@ def makebasis(pos, mass, basis_model, config=None, basis_id='sphereSL', time=0,
 
     if os.path.isfile(modelname) == False:
         print("-> File model not found so we are computing one \n")
-        if empirical == True:
+        if basis_model == "empirical":
             print('-> Computing empirical model')
             rad, rho = empirical_density_profile(pos, mass, nbins=numr)
-        elif empirical == False:
-            makemodel.hernquist_halo()
-
-        makemodel.Profiles(density_profile)
-
-        R, D, M, P = makemodel.makemodel(rho, rvals=np.logspace(np.log10(rmin),
-            np.log10(rmax), numr), M=np.sum(mass), outfile=modelname, return_values=True)
+            R, D, M, P = makemodel('empirical', func=None, dvals=rho, rvals=np.logspace(np.log10(rmin),  np.log10(rmax), numr), M=np.sum(mass), outfile=modelname, return_values=True)
+        elif empirical == "Hernquist":
             print('-> Computing analytical Hernquist model')
-            R, D, M, P = makemodel.makemodel(hernquist_halo, 1, [scale], rvals=, pfile=modelname)
+            #makemodel.hernquist_halo()
+            #R, D, M, P = makemodel.makemodel(hernquist_halo, 1, [scale], rvals=, pfile=modelname)
+
         print('-> Model computed: rmin={}, rmax={}, numr={}'.format(R[0], R[-1], len(R)))
     else:
         R, D, M, P  = np.loadtxt(modelname, skiprows=3, unpack=True) 

EXPtools/basis_builder/makemodel.py

@@ -1,5 +1,47 @@
 import numpy as np
 
+def empirical_density_profile(pos, mass, nbins=500, rmin=0, rmax=600, log_space=False):
+    """
+    Computes the number density radial profile assuming all particles have the same mass.
+
+    Args:
+        pos (ndarray): array of particle positions in cartesian coordinates with shape (n,3).
+        mass (ndarray): array of particle masses with shape (n,).
+        nbins (int, optional): number of bins in the radial profile. Default is 500.
+        rmin (float, optional): minimum radius of the radial profile. Default is 0.
+        rmax (float, optional): maximum radius of the radial profile. Default is 600.
+        log_space (bool, optional): whether to use logarithmic binning. Default is False.
+
+    Returns:
+        tuple: a tuple containing the arrays of radius and density with shapes (nbins,) and (nbins,), respectively.
+
+    Raises:
+        ValueError: if pos and mass arrays have different lengths or if nbins is not a positive integer.
+    """
+    if len(pos) != len(mass):
+        raise ValueError("pos and mass arrays must have the same length")
+    if not isinstance(nbins, int) or nbins <= 0:
+        raise ValueError("nbins must be a positive integer")
+
+    # Compute radial distances
+    r_p = np.sqrt(np.sum(pos**2, axis=1))
+
+    # Compute bin edges and shell volumes
+    if log_space:
+        bins = np.logspace(np.log10(rmin), np.log10(rmax), nbins+1)
+    else:
+        bins = np.linspace(rmin, rmax, nbins+1)
+    V_shells = 4/3 * np.pi * (bins[1:]**3 - bins[:-1]**3)
+
+    # Compute density profile
+    density, _ = np.histogram(r_p, bins=bins, weights=mass)
+    density /= V_shells
+
+    # Compute bin centers and return profile
+    radius = 0.5 * (bins[1:] + bins[:-1])
+
+    return radius, density
+
 def makemodel(basis_type, func, dvals, rvals, M, 
         unit_physical=False, return_values=False,
         outfile='', plabel = '', verbose=True):
@@ -43,9 +85,8 @@ def makemodel(basis_type, func, dvals, rvals, M,
 
     # query out the density values
 
-    if basis_type == 'empirical':
-        dvals = func(rvals)
 
+    #vals = empirical_density_profile(rvals)
 
     # make the mass and potential arrays
     mvals = np.zeros(dvals.size)
@@ -60,13 +101,9 @@ def makemodel(basis_type, func, dvals, rvals, M,
 
     # evaluate mass enclosed and potential energy by recursion
     for indx in range(1, dvals.size):
-        mvals[indx] = mvals[indx-1] 
-                    + 2.0*np.pi*(rvals[indx-1]*rvals[indx-1]*dvals[indx-1] 
-                                 + rvals[indx]*rvals[indx]*dvals[indx])*(rvals[indx] - rvals[indx-1])
+        mvals[indx] = mvals[indx-1] + 2.0*np.pi*(rvals[indx-1]*rvals[indx-1]*dvals[indx-1] + rvals[indx]*rvals[indx]*dvals[indx])*(rvals[indx] - rvals[indx-1])
 
-        pwvals[indx] = pwvals[indx-1] 
-                     + 2.0*np.pi*(rvals[indx-1]*dvals[indx-1] 
-                                 + rvals[indx]*dvals[indx])*(rvals[indx] - rvals[indx-1]);
+        pwvals[indx] = pwvals[indx-1] + 2.0*np.pi*(rvals[indx-1]*dvals[indx-1] + rvals[indx]*dvals[indx])*(rvals[indx] - rvals[indx-1]);
 
     # evaluate potential (see theory document)
     pvals = -mvals/(rvals+1.e-10) - (pwvals[dvals.size-1] - pwvals)

EXPtools/utils/coefficients.py

@@ -0,0 +1,20 @@
+def remove_terms(original_coefficients, n, l, m):
+    """
+    Remove coefficients 
+    """
+
+    assert len(l) == len(m) == len(n)
+    copy_coeffcients = original_coefficients.deepcopy()
+    coefs_matrix = original_coefficients.getAllCoefs()
+    t_snaps = original_coefficients.Times()
+
+
+    for t in range(len(t_snaps)):
+        for i in range(len(l)):
+            lm_idx = int(l[i]*(l[i]+1) / 2) + m[i]
+            print(n[i],l[i],m[i],lm_idx)
+            try: coefs_matrix[n[i], lm_idx, t] = np.complex128(0)
+            except IndexError: continue
+        copy_coeffcients.setMatrix(mat=coefs_matrix[:,:, t], time=t_snaps[t])
+
+    return copy_coeffcients

EXPtools/visuals/__init__.py

@@ -5,4 +5,4 @@
 from .visualize import return_fields_in_grid
 from .visualize import slice_fields
 from .visualize import slice_3d_fields
-
+from .visuals3D import field3Drender

EXPtools/visuals/visualize.py

@@ -1,7 +1,6 @@
 import os,  sys, pickle, pyEXP
 import numpy as np
 
-###Field computations for plotting###
 def make_basis_plot(basis, savefile=None, nsnap='mean', y=0.92, dpi=200):
     """
     Plots the potential of the basis functions for different values of l and n.
@@ -204,7 +203,130 @@ def return_fields_in_grid(basis, coefficients, times=[0],
         elif projection == 'YZ':
             pmin  = [proj_plane,  -grid_lim, -grid_lim]
             pmax  = [proj_plane, grid_lim, grid_lim]
-            grid  = [0, npoints, npoints]
-
-        field_gen = pyEXP.field.FieldGenerator(times, pmin, pmax, grid)
-        return field_gen.slices(basis, coefficients), xgrid
+
+
+    field_gen = pyEXP.field.FieldGenerator(times, pmin, pmax, grid)
+
+    return field_gen.volumes(basis, coefficients), xgrid
+
+def slice_fields(basis, coefficients, time=0, 
+                 projection='XY', proj_plane=0, npoints=300, 
+                 grid_limits=(-300, 300), prop='dens', monopole_only=False):
+    """
+    Plots a slice projection of the fields of a simulation.
+
+    Args:
+    basis (obj): object containing the basis functions for the simulation
+    coefficients (obj): object containing the coefficients for the simulation
+    time (float): the time at which to plot the fields
+    projection (str): the slice projection to plot. Can be 'XY', 'XZ', or 'YZ'.
+    proj_plane (float, optional): the value of the coordinate that is held constant in the slice projection
+    npoints (int, optional): the number of grid points in each dimension
+    grid_limits (tuple, optional): the limits of the grid in the x and y dimensions, in the form (x_min, x_max)
+    prop (str, optional): the property to return. Can be 'dens' (density), 'pot' (potential), or 'force' (force).
+    monopole_only (bool, optional): whether to return the monopole component in the returned property value.
+
+    Returns:
+    array or list: the property specified by `prop`. If `prop` is 'force', a list of the x, y, and z components of the force is returned.
+                    Also returns the grid used to compute slice fields. 
+    """
+    x = np.linspace(grid_limits[0], grid_limits[1], npoints)
+    xgrid = np.meshgrid(x, x)
+    xg = xgrid[0].flatten()
+    yg = xgrid[1].flatten()
+
+
+    if projection not in ['XY', 'XZ', 'YZ']:
+        raise ValueError("Invalid projection specified. Possible values are 'XY', 'XZ', and 'YZ'.")
+
+    N = len(xg)
+    rho0 = np.zeros_like(xg)
+    pot0 = np.zeros_like(xg)
+    rho = np.zeros_like(xg)
+    pot = np.zeros_like(xg)
+    basis.set_coefs(coefficients.getCoefStruct(time))
+
+    for k in range(0, N):
+        if projection == 'XY':
+            rho0[k], pot0[k], rho[k], pot[k], fx, fy, fz = basis.getFields(xg[k], yg[k], proj_plane)
+        elif projection == 'XZ':
+            rho0[k], pot0[k], rho[k], pot[k], fx, fy, fz = basis.getFields(xg[k], proj_plane, yg[k])
+        elif projection == 'YZ':
+            rho0[k], pot0[k], rho[k], pot[k], fx, fy, fz = basis.getFields(proj_plane, xg[k], yg[k])
+
+    dens = rho.reshape(npoints, npoints)
+    pot = pot.reshape(npoints, npoints)
+    dens0 = rho0.reshape(npoints, npoints)
+    pot0 = pot0.reshape(npoints, npoints)
+
+    if prop == 'dens':
+        if monopole_only:
+            return dens0
+        return dens0, dens, xgrid
+
+    if prop == 'pot':
+        if monopole_only:
+            return pot0
+        return pot0, pot, xgrid
+
+    if prop == 'force':
+        return [fx.reshape(npoints, npoints), fy.reshape(npoints, npoints), fz.reshape(npoints, npoints)], xgrid
+
+
+def slice_3d_fields(basis, coefficients, time=0,  npoints=50, 
+                 grid_limits=(-300, 300), prop='dens', monopole_only=False):
+    """
+    Plots a slice projection of the fields of a simulation.
+
+    Args:
+    basis (obj): object containing the basis functions for the simulation
+    coefficients (obj): object containing the coefficients for the simulation
+    time (float): the time at which to plot the fields
+    npoints (int, optional): the number of grid points in each dimension
+    grid_limits (tuple, optional): the limits of the grid in the x and y dimensions, in the form (x_min, x_max)
+    prop (str, optional): the property to return. Can be 'dens' (density), 'pot' (potential), or 'force' (force).
+    monopole_only (bool, optional): whether to return the monopole component in the returned property value.
+
+    Returns:
+    array or list: the property specified by `prop`. If `prop` is 'force', a list of the x, y, and z components of the force is returned.
+                    Also returns the grid used to compute slice fields. 
+    """
+    x = np.linspace(grid_limits[0], grid_limits[1], npoints)
+    xgrid = np.meshgrid(x, x, x)
+    xg = xgrid[0].flatten()
+    yg = xgrid[1].flatten() 
+    zg = xgrid[2].flatten()
+
+
+
+    N = len(xg)
+    rho0 = np.zeros_like(xg)
+    pot0 = np.zeros_like(xg)
+    rho = np.zeros_like(xg)
+    pot = np.zeros_like(xg)
+    fx = np.zeros_like(xg)
+    fy = np.zeros_like(xg)
+    fz = np.zeros_like(xg)
+    basis.set_coefs(coefficients.getCoefStruct(time))
+
+    for k in range(0, N):
+        rho0[k], pot0[k], rho[k], pot[k], fx[k], fy[k], fz[k] = basis.getFields(xg[k], yg[k], zg[k])
+
+    dens = rho.reshape(npoints, npoints, npoints)
+    pot = pot.reshape(npoints, npoints, npoints)
+    dens0 = rho0.reshape(npoints, npoints, npoints)
+    pot0 = pot0.reshape(npoints, npoints, npoints)
+
+    if prop == 'dens':
+        if monopole_only:
+            return dens0
+        return dens0, dens, xgrid
+
+    if prop == 'pot':
+        if monopole_only:
+            return pot0
+        return pot0, pot, xgrid
+
+    if prop == 'force':
+        return [fx.reshape(npoints, npoints, npoints), fy.reshape(npoints, npoints, npoints), fz.reshape(npoints, npoints, npoints)], xgrid
+

EXPtools/visuals/visuals3D.py

@@ -0,0 +1,108 @@
+
+import numpy as np
+import k3d
+
+def field3Dcontour(field, contour_range, contour_name, size, **kwargs):
+    """
+    
+    field: numpy.ndarray
+        shape(t, N, N, N), where t is the time axis, N are the spatial axis.
+    
+    contour_ranges: list
+        values in percentage of the contour levels. e.g [0.5, 0.75] means
+        contours at the 50% and 75% level.  
+    
+    volume_names 
+    
+    kwargs:
+        color map
+    
+    """
+
+    field_max = np.max(np.abs(field))
+    #field_min = np.min(field)
+
+
+    volume =  k3d.volume(field.astype(np.float32), 
+                        alpha_coef=1.0,
+                        color_range=contour_range, #*field_max, 
+                        color_map=(np.array(k3d.colormaps.paraview_color_maps.Cool_to_Warm_Extended).reshape(-1,4) 
+                   * np.array([1,1.0,1.0,1.0])).astype(np.float32),
+                        name=contour_name)
+
+    # Where this values come from? 
+    volume.opacity_function  = [0.        , 0.        , 0.21327923, 0.98025   , 0.32439035,
+           0.        , 0.5       , 0.        , 0.67560965, 0.        ,
+           0.74537706, 0.9915    , 1.        , 0.        ]
+
+
+    volume.transform.bounds = [-size[0], size[0],
+                               -size[1], size[1],
+                               -size[2], size[2]]
+
+    if 'alpha' in kwargs.keys():
+        volume.alpha_coef = kwargs['alpha']
+
+    return volume
+
+def orbit3D(orbit, orbit_name, **kwargs):
+    """
+    # TODO: Add time-dependent 
+    
+    orbit: numpy.array
+        (3, N) shape 
+    
+    orbit_name 
+    
+    color 
+    
+    kwargs:
+        color map
+    
+    """
+
+    orbit_trayectory = k3d.line(orbit.astype(np.float32),
+                                width=1,
+                                color_range=[0, 0.1], color=0x3e3a3a, name=orbit_name)
+
+
+
+    if 'alpha' in kwargs.keys():
+        orbit_trayectory.alpha_coef = kwargs['alpha']
+
+    return orbit_trayectory
+
+
+
+def field3Drender(volumes, contour_ranges, size, contour_names=['Inner', 'Outter'],
+        contour_alphas=[8,1], **kwargs):
+    """
+    volumes: numpy.ndarray
+        shape: t, gridx, gridy, gridz
+        
+    orbit
+
+    """
+    render = k3d.plot(height=800)
+
+    for vol in range(len(contour_ranges)):
+        v = field3Dcontour(volumes[0], 
+                           contour_range=contour_ranges[vol], 
+                           contour_name=contour_names[vol], 
+                           alpha=contour_alphas[vol], size=size)
+        render += v
+        v_t = {}
+        for t in range(volumes.shape[0]):
+            v_t[str(float(t))] = volumes[t].astype(np.float32)
+        v.volume = v_t
+
+    # Implement 3d orbit
+    if 'orbits' in kwargs.keys():
+
+        for orb in range(len(kwargs['orbits_names'])):
+            o = orbit3D(kwargs['orbits'][orb], 
+                        orbit_name=kwargs['orbits_names'][orb], width=10)
+            render += o
+
+
+    return render