allow propagating kwargs

solpolpy/core.py

@@ -14,7 +14,10 @@
 from solpolpy.polarizers import npol_to_mzp
 
 
-def resolve(input_data: t.Union[t.List[str], NDCollection], out_system: str, imax_effect: bool = False) -> NDCollection:
+def resolve(input_data: t.Union[t.List[str], NDCollection],
+            out_system: str,
+            imax_effect: bool = False,
+            out_angles: t.Optional[t.List[float]] = None) -> NDCollection:
     """
     Apply - apply a polarization transformation to a set of input
     dataframes.
@@ -74,14 +77,15 @@ def resolve(input_data: t.Union[t.List[str], NDCollection], out_system: str, ima
     requires_alpha = check_alpha_requirement(transform_path)
 
     if imax_effect:
-        if (input_kind == 'MZP') and (out_system == 'MZP'):
+        if input_kind == 'MZP' and out_system == 'MZP':
             input_data = resolve_imax_effect(input_data)
         else:
             raise UnsupportedTransformationError('IMAX effect applies only for MZP->MZP solpolpy transformations')
 
     if requires_alpha and "alpha" not in input_key:
         input_data = add_alpha(input_data)
-    result = equation(input_data)
+
+    result = equation(input_data, out_angles=out_angles)
 
     return result
 
@@ -171,13 +175,13 @@ def check_alpha_requirement(path: t.List[str]) -> bool:
     return requires_alpha
 
 
-def generate_imax_matrix(arrayshape) -> np.ndarray:
+def generate_imax_matrix(array_shape) -> np.ndarray:
     """
     Define an A matrix with which to convert MZP^ (camera coords) = A x MZP (solar coords)
 
     Parameters
     -------
-    arrayshape
+    array_shape
         Defined input WCS array shape for matrix generation
 
     Returns
@@ -190,7 +194,7 @@ def generate_imax_matrix(arrayshape) -> np.ndarray:
     # Ideal MZP wrt Solar North
     thmzp = [-60, 0, 60] * u.degree
 
-    long_arr, lat_arr = np.meshgrid(np.linspace(-20, 20, arrayshape[0]), np.linspace(-20, 20, arrayshape[1]))
+    long_arr, lat_arr = np.meshgrid(np.linspace(-20, 20, array_shape[0]), np.linspace(-20, 20, array_shape[1]))
 
     # Foreshortening (IMAX) effect on polarizer angle
     phi_m = np.arctan2(np.tan(thmzp[0]) * np.cos(long_arr * u.degree), np.cos(lat_arr * u.degree)).to(u.degree)
@@ -200,7 +204,7 @@ def generate_imax_matrix(arrayshape) -> np.ndarray:
     phi = np.stack([phi_m, phi_z, phi_p])
 
     # Define the A matrix
-    mat_a = np.empty((arrayshape[0], arrayshape[1], 3, 3))
+    mat_a = np.empty((array_shape[0], array_shape[1], 3, 3))
 
     for i in range(3):
         for j in range(3):
@@ -309,10 +313,10 @@ def _compose2(f: t.Callable, g: t.Callable) -> t.Callable:
     Callable
         composed function
     """
-    return lambda *a, **kw: f(g(*a, **kw))
+    return lambda *a, **kw: f(g(*a, **kw), **kw)
 
 
-def identity(x: t.Any) -> t.Any:
+def identity(x: t.Any, **kwargs) -> t.Any:
     """Identity function that returns the input
 
     Parameters

solpolpy/polarizers.py

@@ -17,7 +17,7 @@ def conv_polar_from_head(input_cube):
 
 
 # TODO: prepare a config file where the reference angle say of STEREO, KCor etc can be set
-def npol_to_mzp(input_collection):
+def npol_to_mzp(input_collection, **kwargs):
     """
     Notes
     ------
@@ -70,7 +70,7 @@ def npol_to_mzp(input_collection):
     return NDCollection(Bmzp_cube, meta={}, aligned_axes="all")
 
 
-def mzp_to_bpb(input_collection):
+def mzp_to_bpb(input_collection, **kwargs):
     """
     Notes
     ------
@@ -109,7 +109,7 @@ def mzp_to_bpb(input_collection):
     return NDCollection(BpB_cube, meta={}, aligned_axes="all")
 
 
-def bpb_to_mzp(input_collection):
+def bpb_to_mzp(input_collection, **kwargs):
     """
     Notes
     ------
@@ -140,7 +140,7 @@ def bpb_to_mzp(input_collection):
     return NDCollection(Bmzp_cube, meta={}, aligned_axes="all")
 
 
-def bpb_to_btbr(input_collection):
+def bpb_to_btbr(input_collection, **kwargs):
     """
     Notes
     ------
@@ -167,7 +167,7 @@ def bpb_to_btbr(input_collection):
     return NDCollection(BtBr_cube, meta={}, aligned_axes="all")
 
 
-def btbr_to_bpb(input_collection):
+def btbr_to_bpb(input_collection, **kwargs):
     """
     Notes
     ------
@@ -194,7 +194,7 @@ def btbr_to_bpb(input_collection):
     return NDCollection(BpB_cube, meta={}, aligned_axes="all")
 
 
-def mzp_to_stokes(input_collection):
+def mzp_to_stokes(input_collection, **kwargs):
     """
     Notes
     ------
@@ -219,7 +219,7 @@ def mzp_to_stokes(input_collection):
     return NDCollection(BStokes_cube, meta={}, aligned_axes="all")
 
 
-def stokes_to_mzp(input_collection):
+def stokes_to_mzp(input_collection, **kwargs):
     """
     Notes
     ------
@@ -249,7 +249,7 @@ def stokes_to_mzp(input_collection):
     return NDCollection(Bmzp_cube, meta={}, aligned_axes="all")
 
 
-def mzp_to_bp3(input_collection):
+def mzp_to_bp3(input_collection, **kwargs):
     """
     Notes
     ------
@@ -291,7 +291,7 @@ def mzp_to_bp3(input_collection):
     return NDCollection(Bp3_cube, meta={}, aligned_axes="all")
 
 
-def bp3_to_mzp(input_collection):
+def bp3_to_mzp(input_collection, **kwargs):
     """
     Notes
     ------
@@ -323,7 +323,7 @@ def bp3_to_mzp(input_collection):
     return NDCollection(Bmzp_cube, meta={}, aligned_axes="all")
 
 
-def btbr_to_mzp(input_collection):
+def btbr_to_mzp(input_collection, **kwargs):
     """
     Notes
     ------
@@ -353,7 +353,7 @@ def btbr_to_mzp(input_collection):
     return NDCollection(Bmzp_cube, meta={}, aligned_axes="all")
 
 
-def bp3_to_bthp(input_collection):
+def bp3_to_bthp(input_collection, **kwargs):
     """
     Notes
     ------
@@ -379,7 +379,7 @@ def bp3_to_bthp(input_collection):
     return NDCollection(Bthp_cube, meta={}, aligned_axes="all")
 
 
-def btbr_to_npol(input_collection, angles):
+def btbr_to_npol(input_collection, out_angles=None, **kwargs):
     """
     Notes
     ------
@@ -389,10 +389,13 @@ def btbr_to_npol(input_collection, angles):
     if "alpha" not in input_collection:
         raise ValueError("missing alpha")
 
+    if out_angles is None:
+        raise ValueError("out_angles must be defined as a list of floats")
+
     alpha = input_collection['alpha'].data * u.radian
     Bt, Br = input_collection['Bt'].data, input_collection['Br'].data
 
-    npol_ang = angles
+    npol_ang = out_angles
     Bnpol = {}
     Bnpol_cube = []
     mask = combine_masks(input_collection["Bt"].mask, input_collection["Br"].mask)
@@ -406,7 +409,7 @@ def btbr_to_npol(input_collection, angles):
     return NDCollection(Bnpol_cube, meta={}, aligned_axes="all")
 
 
-def fourpol_to_stokes(input_collection):
+def fourpol_to_stokes(input_collection, **kwargs):
     """
     Notes
     ------

tests/test_core.py

@@ -279,6 +279,11 @@ def test_bp3_to_bthp(bp3_ones):
     assert isinstance(result, NDCollection)
 
 
+def test_btbr_to_npol(btbr_ones):
+    result = resolve(btbr_ones, "npol", out_angles=[5.1, 30.5, 80.5])
+    assert isinstance(result, NDCollection)
+
+
 def test_imax_effect(mzp_data):
     result = resolve(mzp_data, "MZP", imax_effect=True)
     assert isinstance(result, NDCollection)