def test_pull_back_vector_pbc():
    # pull back vector field with periodic boundary conditions
    x = np.linspace(0, 2 * np.pi, 151)
    y = np.linspace(0, 2 * np.pi, 151)
    X, Y = np.meshgrid(x, y)
    # take a torus and simple flow wrapping around twice
    u = np.stack([np.ones(X.shape), np.ones(Y.shape)], axis=-1)
    t_eval = np.linspace(0, 2 * np.pi, 21)
    phi, fold_pbc = flow.rect_flow(u,
                                   x,
                                   y,
                                   None, (t_eval[0], t_eval[-1]),
                                   pbc=(True, True),
                                   t_eval=t_eval)
    vf = np.stack([np.sin(X), np.zeros(Y.shape)], axis=-1)
    vf_pullback_2pi = flow.pull_back(phi,
                                     t_eval,
                                     vf,
                                     x,
                                     y,
                                     fold_pbc=fold_pbc,
                                     t_field=t_eval[20])
    vf_pullback_pi = flow.pull_back(phi,
                                    t_eval,
                                    vf,
                                    x,
                                    y,
                                    fold_pbc=fold_pbc,
                                    t_field=t_eval[10])
    err_2pi = np.linalg.norm(vf_pullback_2pi - vf, axis=-1).flatten()
    err_pi = np.linalg.norm(vf_pullback_pi + vf, axis=-1).flatten()

    assert np.allclose(err_2pi, 0, atol=1e-4) and np.allclose(
        err_pi, 0, atol=1e-4)
def test_pull_back_covector():
    # pull back covector field under hyperbolic flow
    x = np.linspace(-3, 3, 151)
    y = np.linspace(-3, 3, 151)
    X, Y = np.meshgrid(x, y)
    u = np.stack([X, -Y], axis=-1)
    one_form = np.stack([np.ones(X.shape), np.ones(Y.shape)], axis=-1)
    # keep it simple.
    t_eval = np.linspace(0, 1, 21)
    phi = flow.rect_flow(u,
                         x,
                         y,
                         None, (t_eval[0], t_eval[-1]),
                         initial_cond_x=x[50:-50],
                         initial_cond_y=y[50:-50],
                         t_eval=t_eval)
    vf_pullback = flow.pull_back(phi,
                                 t_eval,
                                 one_form,
                                 x,
                                 y,
                                 t_field=t_eval[-1],
                                 covariant=False)
    one_form_pullback = flow.pull_back(phi,
                                       t_eval,
                                       one_form,
                                       x,
                                       y,
                                       t_field=t_eval[-1],
                                       covariant=True)
    # should be a pure shear transformation.

    assert (np.all(np.abs(one_form_pullback[:, :, 1] - np.exp(-1)) <= 0.005)
            and np.all(np.abs(vf_pullback[:, :, 1] - np.exp(1)) <= 0.005))
def test_pull_back_vector_field_rotation():
    # pull back vector field under rotation
    x = np.linspace(-3, 3, 101)
    y = np.linspace(-3, 3, 101)
    X, Y = np.meshgrid(x, y)
    r = np.sqrt(X**2 + Y**2)
    mask = (r <= 2.75).astype(int)
    mask2 = (r <= 2.5).astype(int)
    angle = np.arctan2(Y, X)
    u = np.stack([mask * r * np.sin(angle), -mask * r * np.cos(angle)],
                 axis=-1)
    # points in clockwise direction
    t_eval = np.linspace(0, np.pi / 2, 21)
    phi = flow.rect_flow(u, x, y, None, (t_eval[0], t_eval[-1]), t_eval=t_eval)

    vf = np.stack([np.ones(X.shape), np.zeros(Y.shape)], axis=-1)
    vf_pullback = flow.pull_back(phi, t_eval, vf, x, y, t_field=t_eval[-1])
    sol = np.stack([np.zeros(X.shape), np.ones(Y.shape)], axis=-1)
    err = (np.linalg.norm(vf_pullback - sol, axis=-1) * mask2)
    # one form pullback should give the same for a rotation
    form_pullback = flow.pull_back(phi,
                                   t_eval,
                                   vf,
                                   x,
                                   y,
                                   t_field=t_eval[-1],
                                   covariant=True)
    err_form = (np.linalg.norm(form_pullback - sol, axis=-1) * mask2)

    assert (np.quantile(err, 0.95) < 0.005) and (np.quantile(err_form, 0.95) <
                                                 0.005)
def test_pull_back_radial_function():
    # radial flow with ring-like function, check also density
    x = np.linspace(-2, 2, 101)
    y = np.linspace(-2, 2, 101)
    X, Y = np.meshgrid(x, y)
    r = np.sqrt(X**2 + Y**2)

    def ring(r1, r2):
        return ((r >= r1) & (r <= r2)).astype(float)

    radial_vf = -np.stack([X, Y], axis=-1)
    t_eval = np.linspace(0, 1, 51)
    phi = flow.rect_flow(radial_vf, x, y, None, (0, 1), t_eval=t_eval)
    f = ring(0.5, 1)
    f_pullbacks = [
        flow.pull_back(phi, t_eval, f, x, y, t_field=t, density=True)
        for t in t_eval
    ]
    sol = [
        ring(np.exp(t) * 0.5,
             np.exp(t) * 1) * np.exp(-2 * t) for t in t_eval
    ]
    # ring moves and is dilutes
    err = [
        np.mean(np.abs(a - b)) / np.mean(b) for a, b, in zip(f_pullbacks, sol)
    ]
    # error is due to numerical issues at ring boundaries.
    assert np.all(np.array(err) < 0.1)
def test_pull_back_trivial():
    # simple translation, but with vectors
    x = np.linspace(-2, 2, 41)
    y = np.linspace(-2, 2, 21)
    X, Y = np.meshgrid(x, y)
    t_eval = np.arange(10)
    phi = np.stack([np.stack([X + step, Y], axis=-1) for step in t_eval])
    f = np.stack([X**2, np.zeros(X.shape)], axis=-1)
    f_pullback = flow.pull_back(phi, t_eval, f, x + 9, 2 * y)
    # y+1/2 raises assertion error as it should
    assert np.allclose(f, f_pullback)
def test_pull_back_matrix():
    # pulling back (2, 0) tensor
    # let's take the example of a disclination  of type -1/2 and a rotation
    # vector field
    x = np.linspace(-3, 3, 101)
    y = np.linspace(-3, 3, 101)
    X, Y = np.meshgrid(x, y)
    r = np.sqrt(X**2 + Y**2)
    mask = (r <= 2.75).astype(int)
    mask2 = (r <= 2.5).astype(int)
    u = np.stack([mask * Y, -mask * X], axis=-1)
    t_eval = np.linspace(0, 2 * np.pi / 3, 21)
    phi = flow.rect_flow(u, x, y, None, (t_eval[0], t_eval[-1]), t_eval=t_eval)
    m = np.stack([[X, -Y], [-Y, -X]]).transpose(2, 3, 0, 1)
    m_pullback = flow.pull_back(phi, t_eval, m, x, y, t_field=t_eval[-1])
    err = np.linalg.norm(m - m_pullback, axis=(2, 3)) * mask2

    assert np.quantile(err, 0.9) < 0.01