move stuff under fdmath

meanas/fdmath/__init__.py

@@ -91,4 +91,12 @@ and
   while \\( \\hat{h} \\) and \\( \\tilde{h} \\) are located at \\((m \pm \\frac{1}{2}, n \\pm \\frac{1}{2}, p \\pm \\frac{1}{2})\\)
   with components at \\((m, n \\pm \\frac{1}{2}, p \\pm \\frac{1}{2})\\) etc.
 
+TODO: Explain fdfield_t vs vfdfield_t  / operators vs functional
+TODO: explain dxes
+
 """
+
+from .types import fdfield_t, vfdfield_t, dx_lists_t, fdfield_updater_t
+from .vectorization import vec, unvec
+from . import operators, functional, types, vectorization
+

meanas/fdmath/functional.py

@@ -6,10 +6,10 @@ Basic discrete calculus etc.
 from typing import List, Callable, Tuple, Dict
 import numpy
 
-from .. import field_t, field_updater
+from .types import fdfield_t, fdfield_updater_t
 
 
-def deriv_forward(dx_e: List[numpy.ndarray] = None) -> field_updater:
+def deriv_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t:
     """
     Utility operators for taking discretized derivatives (backward variant).
 
@@ -31,7 +31,7 @@ def deriv_forward(dx_e: List[numpy.ndarray] = None) -> field_updater:
     return derivs
 
 
-def deriv_back(dx_h: List[numpy.ndarray] = None) -> field_updater:
+def deriv_back(dx_h: List[numpy.ndarray] = None) -> fdfield_updater_t:
     """
     Utility operators for taking discretized derivatives (forward variant).
 
@@ -53,7 +53,7 @@ def deriv_back(dx_h: List[numpy.ndarray] = None) -> field_updater:
     return derivs
 
 
-def curl_forward(dx_e: List[numpy.ndarray] = None) -> field_updater:
+def curl_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t:
     """
     Curl operator for use with the E field.
 
@@ -67,7 +67,7 @@ def curl_forward(dx_e: List[numpy.ndarray] = None) -> field_updater:
     """
     Dx, Dy, Dz = deriv_forward(dx_e)
 
-    def ce_fun(e: field_t) -> field_t:
+    def ce_fun(e: fdfield_t) -> fdfield_t:
         output = numpy.empty_like(e)
         output[0] = Dy(e[2])
         output[1] = Dz(e[0])
@@ -80,7 +80,7 @@ def curl_forward(dx_e: List[numpy.ndarray] = None) -> field_updater:
     return ce_fun
 
 
-def curl_back(dx_h: List[numpy.ndarray] = None) -> field_updater:
+def curl_back(dx_h: List[numpy.ndarray] = None) -> fdfield_updater_t:
     """
     Create a function which takes the backward curl of a field.
 
@@ -94,7 +94,7 @@ def curl_back(dx_h: List[numpy.ndarray] = None) -> field_updater:
     """
     Dx, Dy, Dz = deriv_back(dx_h)
 
-    def ch_fun(h: field_t) -> field_t:
+    def ch_fun(h: fdfield_t) -> fdfield_t:
         output = numpy.empty_like(h)
         output[0] = Dy(h[2])
         output[1] = Dz(h[0])

meanas/fdmath/operators.py

@@ -7,7 +7,7 @@ from typing import List, Callable, Tuple, Dict
 import numpy
 import scipy.sparse as sparse
 
-from .. import field_t, vfield_t
+from .types import fdfield_t, vfdfield_t
 
 
 def rotation(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmatrix:
@@ -155,7 +155,7 @@ def cross(B: List[sparse.spmatrix]) -> sparse.spmatrix:
                         [-B[1], B[0], zero]])
 
 
-def vec_cross(b: vfield_t) -> sparse.spmatrix:
+def vec_cross(b: vfdfield_t) -> sparse.spmatrix:
     """
     Vector cross product operator
 

meanas/fdmath/types.py

@@ -0,0 +1,40 @@
+"""
+Types shared across multiple submodules
+"""
+import numpy
+from typing import List, Callable
+
+
+# Field types
+# TODO: figure out a better way to set the docstrings without creating actual subclasses?
+#   Probably not a big issue since they're only used for type hinting
+class fdfield_t(numpy.ndarray):
+    """
+    Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)
+
+    This is actually is just an unaltered `numpy.ndarray`
+    """
+    pass
+
+class vfdfield_t(numpy.ndarray):
+    """
+    Linearized vector field (single vector of length 3*X*Y*Z)
+
+    This is actually just an unaltered `numpy.ndarray`
+    """
+    pass
+
+
+dx_lists_t = List[List[numpy.ndarray]]
+'''
+ 'dxes' datastructure which contains grid cell width information in the following format:
+
+     [[[dx_e_0, dx_e_1, ...], [dy_e_0, ...], [dz_e_0, ...]],
+      [[dx_h_0, dx_h_1, ...], [dy_h_0, ...], [dz_h_0, ...]]]
+
+   where `dx_e_0` is the x-width of the `x=0` cells, as used when calculating dE/dx,
+   and `dy_h_0` is  the y-width of the `y=0` cells, as used when calculating dH/dy, etc.
+'''
+
+
+fdfield_updater_t = Callable[[fdfield_t], fdfield_t]

meanas/fdmath/vectorization.py

@@ -0,0 +1,47 @@
+"""
+Functions for moving between a vector field (list of 3 ndarrays, [f_x, f_y, f_z])
+and a 1D array representation of that field [f_x0, f_x1, f_x2,... f_y0,... f_z0,...].
+Vectorized versions of the field use row-major (ie., C-style) ordering.
+"""
+
+from typing import List
+import numpy
+
+from .types import fdfield_t, vfdfield_t
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def vec(f: fdfield_t) -> vfdfield_t:
+    """
+    Create a 1D ndarray from a 3D vector field which spans a 1-3D region.
+
+    Returns None if called with f=None.
+
+    :param f: A vector field, [f_x, f_y, f_z] where each f_ component is a 1 to
+        3D ndarray (f_* should all be the same size). Doesn't fail with f=None.
+    :return: A 1D ndarray containing the linearized field (or None)
+    """
+    if numpy.any(numpy.equal(f, None)):
+        return None
+    return numpy.ravel(f, order='C')
+
+
+def unvec(v: vfdfield_t, shape: numpy.ndarray) -> fdfield_t:
+    """
+    Perform the inverse of vec(): take a 1D ndarray and output a 3D field
+     of form [f_x, f_y, f_z] where each of f_* is a len(shape)-dimensional
+     ndarray.
+
+    Returns None if called with v=None.
+
+    :param v: 1D ndarray representing a 3D vector field of shape shape (or None)
+    :param shape: shape of the vector field
+    :return: [f_x, f_y, f_z] where each f_ is a len(shape) dimensional ndarray
+     (or None)
+    """
+    if numpy.any(numpy.equal(v, None)):
+        return None
+    return v.reshape((3, *shape), order='C')
+