From 429e63e0d9904fb5180c3c853d448dbe34672d51 Mon Sep 17 00:00:00 2001
From: Joern Ungermann <j.ungermann@fz-juelich.de>
Date: Fri, 2 Dec 2022 15:43:05 +0100
Subject: [PATCH] Added 2-D regularization matrices and simplified
 block-matrix-code.

---
 jutil/_test/test_regularization.py | 101 ++++++++-
 jutil/regularization.py            | 320 ++++++++++++++++++++++-------
 2 files changed, 349 insertions(+), 72 deletions(-)

diff --git a/jutil/_test/test_regularization.py b/jutil/_test/test_regularization.py
index b8c74ba..cade70c 100644
--- a/jutil/_test/test_regularization.py
+++ b/jutil/_test/test_regularization.py
@@ -27,7 +27,25 @@ import jutil.regularization as reg
      (reg.create_l2, np.pi, 8 * np.pi, 4.12855, 33.0384)
      ])
 def test_l(xs, rel, func, val0, val1, val2, val3):
-    L = func(xs, np.ones_like(xs), 1)
+    L = func(xs, 1, 1)
+
+    ys = np.sin(xs)
+    temp = L.dot(ys)
+    # int_0^2pi sin(x)^2 dx = pi
+    # int_0^2pi cos(x)^2 dx = pi
+    # int_0^2pi (-sin(x))^2 dx = pi
+    assert temp.dot(temp).sum() == pytest.approx(val0, rel=rel)
+
+    L = func(xs, 2, 1)
+
+    ys = np.sin(xs)
+    temp = L.dot(ys)
+    # int_0^2pi sin(x)^2 dx = pi
+    # int_0^2pi cos(x)^2 dx = pi
+    # int_0^2pi (-sin(x))^2 dx = pi
+    assert temp.dot(temp).sum() == pytest.approx(0.25 * val0, rel=rel)
+
+    L = func(xs, np.ones_like(xs), np.ones(len(xs) - 1))
 
     ys = np.sin(xs)
     temp = L.dot(ys)
@@ -102,8 +120,10 @@ def test_generate_regularization():
     scales0 = [10., 10.]
     scales1 = [5., 10.]
     Ls = reg.generate_regularization(axes, scales0=scales0, scales1=scales1)
-    assert np.allclose(L0_a, Ls[0].toarray())
-    assert np.allclose(L1_a, Ls[1].toarray())
+    assert np.allclose(L0_a[:3], Ls.toarray()[:3])
+    assert np.allclose(L0_a[3:5], Ls.toarray()[5:7])
+    assert np.allclose(L1_a[:2], Ls.toarray()[3:5])
+    assert np.allclose(L1_a[2:], Ls.toarray()[7:])
 
     Sa_inv_a = np.array([[1.5, -1, 0, 0, 0, 0],
                          [-1, 1.5, 0, 0, 0, 0],
@@ -119,3 +139,78 @@ def test_generate_regularization():
     Ls = reg.generate_regularization(axes, list_stds=list_stds, list_corr=list_corr)
     Sa_inv_t = reg.generate_inverse_covmatrix(Ls)
     assert np.allclose(Sa_inv_a, Sa_inv_t.toarray())
+
+
+def test2d_l0():
+    corr_x, corr_z = 1, 1
+    val0 = 9.8696
+    val1 = 20242.2
+    rel = 0.01
+    xs = np.linspace(0, 2 * np.pi, 21)
+    zs = np.linspace(0, 1 * np.pi, 20)
+    xxs, zzs = np.meshgrid(xs, zs)
+    L0 = reg.create2d_l0(xs, zs, np.ones_like(xxs), corr_x, corr_z)
+
+    ys = np.sin(xxs + zzs).reshape(-1)
+    temp = L0.dot(ys)
+    # int_0^2pi sin(x)^2 dx
+    # int_0^2pi cos(x)^2 dx
+    # int_0^2pi (-sin(x))^2 dx
+    assert pytest.approx(val0, rel=rel) == temp.dot(temp).sum()
+
+    ys = (xxs ** 2 * zzs).reshape(-1)
+    temp = L0.dot(ys)
+    # int_0^2pi x^4 dx
+    # int_0^2pi (2x)^2 dx
+    # int_0^2pi 4 dx
+    assert pytest.approx(val1, rel=rel) == temp.dot(temp).sum()
+
+
+def test2d_l1x():
+    corr_x, corr_z = 1, 1
+    val0 = 9.8696
+    val1 = 3418
+    rel = 0.01
+    xs = np.linspace(0, 2 * np.pi, 21)
+    zs = np.linspace(0, 1 * np.pi, 20)
+    xxs, zzs = np.meshgrid(xs, zs)
+    L0 = reg.create2d_l1x(xs, zs, np.ones_like(xxs), corr_x, corr_z)
+
+    ys = np.sin(xxs + zzs).reshape(-1)
+    temp = L0.dot(ys)
+    # int_0^2pi sin(x)^2 dx
+    # int_0^2pi cos(x)^2 dx
+    # int_0^2pi (-sin(x))^2 dx
+    assert pytest.approx(val0, rel=rel) == temp.dot(temp).sum()
+
+    ys = (xxs ** 2 * zzs).reshape(-1)
+    temp = L0.dot(ys)
+    # int_0^2pi x^4 dx
+    # int_0^2pi (2x)^2 dx
+    # int_0^2pi 4 dx
+    assert pytest.approx(val1, rel=rel) == temp.dot(temp).sum()
+
+
+def test2d_l1y():
+    corr_x, corr_z = 1, 1
+    val0 = 9.8696
+    val1 = 6152.89
+    rel = 0.01
+    xs = np.linspace(0, 2 * np.pi, 21)
+    zs = np.linspace(0, 1 * np.pi, 20)
+    xxs, zzs = np.meshgrid(xs, zs)
+    L0 = reg.create2d_l1y(xs, zs, np.ones_like(xxs), corr_x, corr_z)
+
+    ys = np.sin(xxs + zzs).reshape(-1)
+    temp = L0.dot(ys)
+    # int_0^2pi sin(x)^2 dx
+    # int_0^2pi cos(x)^2 dx
+    # int_0^2pi (-sin(x))^2 dx
+    assert pytest.approx(val0, rel=rel) == temp.dot(temp).sum()
+
+    ys = (xxs ** 2 * zzs).reshape(-1)
+    temp = L0.dot(ys)
+    # int_0^2pi x^4 dx
+    # int_0^2pi (2x)^2 dx
+    # int_0^2pi 4 dx
+    assert pytest.approx(val1, rel=rel) == temp.dot(temp).sum()
diff --git a/jutil/regularization.py b/jutil/regularization.py
index a59ead9..de9b4d4 100644
--- a/jutil/regularization.py
+++ b/jutil/regularization.py
@@ -1,7 +1,29 @@
+import itertools
 import numpy as np
 import scipy.sparse as sp
 
 
+def _compute_integration_weights(dxs):
+    """ Helper function to compute weights from the trapezoid rule from
+        interval length.
+
+    Parameters
+    ----------
+    dxs : 1darray (N,)
+        lengths of intervals
+
+    Returns
+    -------
+    1darray
+        weights
+
+    """
+    dx_int = np.zeros(len(dxs) + 1)
+    dx_int[:-1] += dxs
+    dx_int[1:] += dxs
+    return dx_int / 2
+
+
 def create_l0(xs, stds, corr):
     """ This function generates the L0 regularization
         in sparse representation.
@@ -39,10 +61,8 @@ def create_l0(xs, stds, corr):
         ISBN 978-0-89871-572-9 978-0-89871-792-1
     """
     dxs = np.diff(xs) / corr
-    dx_int = np.zeros_like(xs)
-    dx_int[:-1] += dxs
-    dx_int[1:] += dxs
-    return sp.diags(np.sqrt(dx_int / 2) / stds)
+    dx_int = _compute_integration_weights(dxs)
+    return sp.diags(np.sqrt(dx_int) / stds)
 
 
 def create_l1(xs, stds, corr):
@@ -83,11 +103,9 @@ def create_l1(xs, stds, corr):
     """
 
     dxs = np.diff(xs) / corr
-    if np.asarray(stds).ndim == 0:
-        stds = np.tile(stds, len(xs))
-    if np.asarray(corr).ndim == 0:
-        corr = np.tile(corr, len(xs))
-    scale = np.sqrt(1 / dxs) / stds[:-1]
+    if isinstance(stds, np.ndarray) and stds.ndim == 1:
+        stds = stds[:-1]
+    scale = np.sqrt(1 / dxs) / stds
     return sp.diags([-scale, scale], offsets=[0, 1],
                     shape=(len(scale), len(xs)))
 
@@ -113,9 +131,6 @@ def create_l2(xs, stds, corr):
 
     Note
     ----
-    -   This approximates the second integral in Eq.
-        (7.311: Tarantola). The first derivative is
-        approximated by forward finite difference.
     -   Caution: varying correlation lengths returns
         errors; best practice to only use L2 with constant
         correlation length;
@@ -128,16 +143,14 @@ def create_l2(xs, stds, corr):
 
     dxs = np.diff(xs) / corr
     dmxs = (dxs[1:] + dxs[:-1]) / 2
-    if np.asarray(stds).ndim == 0:
-        stds = np.tile(stds, len(xs))
-    if np.asarray(corr).ndim == 0:
-        corr = np.tile(corr, len(xs))
+    if isinstance(stds, np.ndarray) and stds.ndim == 1:
+        stds = stds[1:-1]
 
     # integration
     dx_int = (dxs[1:] + dxs[:-1])
     dx_int[0] += dxs[0]
     dx_int[-1] += dxs[-1]
-    dx_int = np.sqrt(dx_int / 2) / stds[1:-1]
+    dx_int = np.sqrt(dx_int / 2) / stds
 
     # second order finite difference differences
     dif0 = dx_int / (dxs[:-1] * dmxs)
@@ -148,7 +161,196 @@ def create_l2(xs, stds, corr):
                     offsets=[0, 1, 2], shape=(len(dx_int), len(xs)))
 
 
-def generate_regblock(xs, scale0=1., scale1=0., scale2=0., stds=1, corr=1):
+def create2d_l0(xs, ys, stds, corr_x, corr_y):
+    """ This function generates the L0 regularization
+        in sparse representation.
+
+    Parameters
+    ----------
+    xs : 1darray (N,)
+        axis of altitude layers
+    ys : 1darray (M,)
+        secondary axis
+    stds : 2darray (M, N) or float
+        standard deviations
+    corr_x : 1darray (N-1,) or float
+        correlation lengths of the spacing between two altitude layers
+    corr_y : 1darray (M-1,) or float
+        correlation lengths of the spacing on secondary axis
+
+    Returns
+    -------
+    2darray
+        sparse L0 regularization matrix
+
+    Note
+    ----
+        This approximates the first integral
+        in Eq. (7.311: Tarantola) by trapezoidal
+        rule with flexible clorrelation length
+        and standard deviation.
+
+    Use as such
+    -----------
+        temp = L.dot(ys)
+        norm = (temp.T.dot(temp)).sum()
+
+    Ref
+    ---
+        Inverse Problem Theory and Methods for
+        Model Parameter Estimation, Tarantola, A (2005),
+        ISBN 978-0-89871-572-9 978-0-89871-792-1
+
+    Note
+    ----
+        The 2d integral is handled by mapping the 2-D state to
+        a 1-D vector.
+        It is then handled as in the 1d case by
+        approximating the first integral
+        in Eq. (7.311: Tarantola) by trapezoidal
+        rule with flexible correlation length
+        and standard deviation.
+
+    """
+    dxs = np.diff(xs) / corr_x
+    dx_int = _compute_integration_weights(dxs)
+    dys = np.diff(ys) / corr_y
+    dy_int = _compute_integration_weights(dys)
+    wgts = dx_int[np.newaxis, :] * dy_int[:, np.newaxis]
+    return sp.diags((np.sqrt(wgts) / stds).reshape(-1))
+
+
+def create2d_l1x(xs, ys, stds, corr_x, corr_y):
+    """
+    Creates a L1 regularization matrix to compute the L2 Norm of the
+    first partial derivative in x-direction of a discretized function.
+
+    Parameters
+    ----------
+    xs : 1darray (N,)
+        axis of altitude layers
+    ys : 1darray (M,)
+        secondary axis
+    stds : 2darray (M, N) or float
+        standard deviations
+    corr_x : 1darray (N-1,) or float
+        correlation lengths of the spacing between two altitude layers
+    corr_y : 1darray (M-1,) or float
+        correlation lengths of the spacing on secondary axis
+
+    Returns
+    -------
+    2darray
+        sparse L1 regularization matrix
+
+    Note
+    ----
+        This approximates the second integral in Eq.
+        (7.311: Tarantola). The first derivative is
+        approximated by forward finite difference.
+
+    Use as such
+    -----------
+        temp = L.dot(ys)
+        norm = (temp.T.dot(temp)).sum()
+
+    Ref:
+    ----
+        Inverse Problem Theory and Methods for
+        Model Parameter Estimation, Tarantola, A (2005),
+        ISBN 978-0-89871-572-9 978-0-89871-792-1
+
+
+    Note
+    ----
+        The 2d integral is handled by mapping the 2-D state to
+        a 1-D vector.
+        It is then handled as in the 1d case by
+        approximating the first integral
+        in Eq. (7.311: Tarantola) by trapezoidal
+        rule with flexible correlation length
+        and standard deviation.
+
+    """
+    assert len(xs) == stds.shape[1] and len(ys) == stds.shape[0]
+    stds = stds.reshape(-1)
+    dxs_sqrt = np.sqrt(np.diff(xs) / corr_x)
+    dys = np.diff(ys) / corr_y
+    dy_int_sqrt = np.sqrt(_compute_integration_weights(dys))
+
+    scale = np.zeros_like(stds)
+    for i in range(len(ys)):
+        sl = slice(i * len(xs), (i + 1) * len(xs) - 1)
+        scale[sl] = dy_int_sqrt[i] / (dxs_sqrt * stds[sl])
+    return sp.diags([-scale, scale], offsets=[0, 1],
+                    shape=(len(stds), len(stds)))
+
+
+def create2d_l1y(xs, ys, stds, corr_x, corr_y):
+    """
+    Creates a L1 regularization matrix to compute the L2 Norm of the
+    first partial derivative in y-direction of a discretized function.
+
+    Parameters
+    ----------
+    xs : 1darray (N,)
+        axis of altitude layers
+    ys : 1darray (M,)
+        secondary axis
+    stds : 2darray (M, N) or float
+        standard deviations
+    corr_x : 1darray (N-1,) or float
+        correlation lengths of the spacing between two altitude layers
+    corr_y : 1darray (M-1,) or float
+        correlation lengths of the spacing on secondary axis
+
+    Returns
+    -------
+    2darray
+        sparse L1 regularization matrix
+
+    Note
+    ----
+        This approximates the second integral in Eq.
+        (7.311: Tarantola). The first derivative is
+        approximated by forward finite difference.
+
+    Use as such
+    -----------
+        temp = L.dot(ys)
+        norm = (temp.T.dot(temp)).sum()
+
+    Ref:
+    ----
+        Inverse Problem Theory and Methods for
+        Model Parameter Estimation, Tarantola, A (2005),
+        ISBN 978-0-89871-572-9 978-0-89871-792-1
+
+    Note
+    ----
+        The 2d integral is handled by mapping the 2-D state to
+        a 1-D vector.
+        It is then handled as in the 1d case by
+        approximating the first integral
+        in Eq. (7.311: Tarantola) by trapezoidal
+        rule with flexible correlation length
+        and standard deviation.
+
+    """
+    assert len(xs) == stds.shape[1] and len(ys) == stds.shape[0]
+    stds = stds.reshape(-1)
+    dxs = np.diff(xs) / corr_x
+    dys_sqrt = np.sqrt(np.diff(ys) / corr_y)
+    dx_int_sqrt = np.sqrt(_compute_integration_weights(dxs))
+    scale = np.zeros_like(stds)
+    for i in range(len(ys) - 1):
+        sl = slice(i * len(xs), (i + 1) * len(xs))
+        scale[sl] = dx_int_sqrt / (dys_sqrt[i] * stds[sl])
+    return sp.diags([-scale, scale], offsets=[0, len(xs)],
+                    shape=(len(stds), len(stds)))
+
+
+def generate_regblock(xs, scale0=1, scale1=0, scale2=0, stds=1, corr=1):
     """ This function generate the inverse covariance matrix
 
     Parameters
@@ -174,25 +376,23 @@ def generate_regblock(xs, scale0=1., scale1=0., scale2=0., stds=1, corr=1):
 
     xs = xs.astype(float)
 
-    n = len(xs)
-    Ls = []
-    if np.allclose(scale0, 0.):
-        Ls.append(sp.coo_matrix((n, n)))
-    else:
-        Ls.append(scale0 * create_l0(xs, stds, corr))
-    if np.allclose(scale1, 0.):
-        Ls.append(sp.coo_matrix((n, n)))
-    else:
-        Ls.append(scale1 * create_l1(xs, stds, corr))
-    if np.allclose(scale2, 0.):
-        Ls.append(sp.coo_matrix((n, n)))
-    else:
-        Ls.append(scale2 * create_l2(xs, stds, corr))
-
-    return Ls
-
-
-def generate_regularization(axes, scales0=None, scales1=None, scales2=None, list_stds=None, list_corr=None):
+    Ls = [
+        scale * create(xs, stds, corr)
+        for create, scale in zip(
+            (create_l0, create_l1, create_l2),
+            (scale0, scale1, scale2))
+        if not np.allclose(scale, 0.)]
+
+    return sp.vstack(Ls)
+
+
+def generate_regularization(
+        axes,
+        scales0=itertools.repeat(1),
+        scales1=itertools.repeat(1),
+        scales2=itertools.repeat(0),
+        list_stds=itertools.repeat(1),
+        list_corr=itertools.repeat(1)):
     """ This function generates blockwise the L-regularisation matrices.
 
     Parameters
@@ -200,44 +400,28 @@ def generate_regularization(axes, scales0=None, scales1=None, scales2=None, list
     axes : list os 1darrays
         containing axes for the different blocks of state variable
     scales0 : list, optional
-        scale for L0 regularization matrix for each block, by default None
+        scale for L0 regularization matrix for each block, by default 1
     scales1 : list, optional
-        scale for L1 regularization matrix for each block, by default None
+        scale for L1 regularization matrix for each block, by default 1
     scales2 : list, optional
-        scale for L2 regularization matrix for each block, by default None
-    list_stds : list os floats or 1darrays
-        standard deviations of state variable for each block; by default None
-    list_corr : list os floats or 1darrays
-        correlation lengths, by default None
+        scale for L2 regularization matrix for each block, by default 0
+    list_stds : list of floats or 1darrays, optional
+        standard deviations of state variable for each block; by default 1
+    list_corr : list of floats or 1darrays, optional
+        correlation lengths, by default 1
 
     Returns
     -------
     Ls : list of sparse 2darrays
         L-matrices
     """
+    Ls = [
+        generate_regblock(axis, scale0=scale0, scale1=scale1, scale2=scale2,
+                          stds=stds, corr=corr)
+        for axis, scale0, scale1, scale2, stds, corr
+        in zip(axes, scales0, scales1, scales2, list_stds, list_corr)]
 
-    if scales0 is None:
-        scales0 = [1.] * len(axes)
-    if scales1 is None:
-        scales1 = [1.] * len(axes)
-    if scales2 is None:
-        scales2 = [0.] * len(axes)
-    if list_stds is None:
-        list_stds = [1] * len(axes)
-    if list_corr is None:
-        list_corr = [1] * len(axes)
-    assert len(axes) == len(scales0) == len(scales1) == len(scales2) == len(list_stds) == len(list_corr)
-    L0, L1, L2 = [], [], []
-    for (axis, scale0, scale1, scale2, stds, corr) in zip(axes, scales0, scales1, scales2, list_stds, list_corr):
-        tmp = generate_regblock(axis, scale0=scale0, scale1=scale1, scale2=scale2, stds=stds)
-        L0.append(tmp[0])
-        L1.append(tmp[1])
-        L2.append(tmp[2])
-
-    L0 = sp.block_diag(L0)
-    L1 = sp.block_diag(L1)
-    L2 = sp.block_diag(L2)
-    return [L0, L1, L2]
+    return sp.block_diag(Ls)
 
 
 def generate_inverse_covmatrix(Ls):
@@ -260,6 +444,4 @@ def generate_inverse_covmatrix(Ls):
     Inverse Problem Theory and Methods for Model Parameter Estimation,
     Tarantola, A (2005), ISBN 978-0-89871-572-9 978-0-89871-792-1
     """
-    Sa_inv = 0.5 * (Ls[0].T.dot(Ls[0]) + Ls[1].T.dot(Ls[1]) + Ls[2].T.dot(Ls[2]))
-
-    return Sa_inv
+    return 0.5 * Ls.T.dot(Ls)
-- 
GitLab