<commit_before>from __future__ import division

from __future__ import absolute_import

from firedrake import *

from firedrake_da import *

import numpy as np

import pytest


def test_coarsening_localisation_single_cell_dg0():

    mesh = UnitIntervalMesh(1)

    mesh_hierarchy = MeshHierarchy(mesh, 3)

    r_loc = 2

    fs_hierarchy = tuple([FunctionSpace(m, 'DG', 0) for m in mesh_hierarchy])

    WLoc = Function(fs_hierarchy[-1]).assign(1.0)
    WLoc_ = Function(fs_hierarchy[-1]).assign(WLoc)

    WLoc = CoarseningLocalisation(WLoc, r_loc)

    assert norm(assemble(WLoc - WLoc_)) == 0


def test_coarsening_localisation_single_cell_dg0_2():

    mesh = UnitIntervalMesh(1)

    mesh_hierarchy = MeshHierarchy(mesh, 1)

    r_loc = 1

    fs_hierarchy = tuple([FunctionSpace(m, 'DG', 0) for m in mesh_hierarchy])

    WLoc = Function(fs_hierarchy[1])
    WLoc.dat.data[0] = 1.0

    WLoc = CoarseningLocalisation(WLoc, r_loc)

    assert np.abs(0.5 - WLoc.dat.data[0]) < 1e-5
    assert np.abs(0.5 - WLoc.dat.data[1]) < 1e-5


def test_coarsening_localisation_no_localisation():

    mesh = UnitIntervalMesh(1)

    mesh_hierarchy = MeshHierarchy(mesh, 3)

    r_loc = 0

    fs_hierarchy = tuple([FunctionSpace(m, 'DG', 0) for m in mesh_hierarchy])

    WLoc = Function(fs_hierarchy[-1])

    WLoc.dat.data[0] += 1.0

    WLoc_ = Function(fs_hierarchy[-1]).assign(WLoc)

    WLoc = CoarseningLocalisation(WLoc, r_loc)

    assert norm(assemble(WLoc - WLoc_)) == 0


def test_coarsening_localisation_no_hierarchy():

    mesh = UnitIntervalMesh(1)

    fs = FunctionSpace(mesh, 'DG', 0)

    r_locs = [0, 1]

    for r in r_locs:

        WLoc = Function(fs)

        WLoc.dat.data[0] += 1.0

        WLoc_ = Function(fs).assign(WLoc)

        if r == 1:
            with pytest.raises(Exception):
                CoarseningLocalisation(WLoc, r)

        if r == 0:
            WLoc = CoarseningLocalisation(WLoc, r)
            assert norm(assemble(WLoc - WLoc_)) == 0


def test_covariance_matrix_1():

    mesh = UnitIntervalMesh(10)

    vfs = VectorFunctionSpace(mesh, 'DG', 0, dim=10)

    r_loc = 2

    L = CovarianceLocalisation(vfs, r_loc)

    assert np.shape(L.dat.data) == (10, 10)

    assert np.max(np.abs(np.diagonal(L.dat.data) - np.ones(10))) < 1e-5
    assert np.max(np.abs(np.diagonal(L.dat.data, 1) - (0.75 * np.ones(9)))) < 1e-5
    assert np.max(np.abs(np.diagonal(L.dat.data, 2) - (0.25 * np.ones(8)))) < 1e-5
    assert np.max(np.abs(np.diagonal(L.dat.data, -1) - (0.75 * np.ones(9)))) < 1e-5
    assert np.max(np.abs(np.diagonal(L.dat.data, -2) - (0.25 * np.ones(8)))) < 1e-5


if __name__ == "__main__":
    import os
    pytest.main(os.path.abspath(__file__))<commit_msg>added 2 more tests for localisation, for 2d and cg1 projection<commit_after>from __future__ import division

from __future__ import absolute_import

from firedrake import *

from firedrake_da import *

import numpy as np

import pytest


def test_coarsening_localisation_single_cell_dg0():

    mesh = UnitIntervalMesh(1)

    mesh_hierarchy = MeshHierarchy(mesh, 3)

    r_loc = 2

    fs_hierarchy = tuple([FunctionSpace(m, 'DG', 0) for m in mesh_hierarchy])

    WLoc = Function(fs_hierarchy[-1]).assign(1.0)
    WLoc_ = Function(fs_hierarchy[-1]).assign(WLoc)

    WLoc = CoarseningLocalisation(WLoc, r_loc)

    assert norm(assemble(WLoc - WLoc_)) == 0


def test_coarsening_localisation_single_cell_dg0_2():

    mesh = UnitIntervalMesh(1)

    mesh_hierarchy = MeshHierarchy(mesh, 1)

    r_loc = 1

    fs_hierarchy = tuple([FunctionSpace(m, 'DG', 0) for m in mesh_hierarchy])

    WLoc = Function(fs_hierarchy[1])
    WLoc.dat.data[0] = 1.0

    WLoc = CoarseningLocalisation(WLoc, r_loc)

    assert np.abs(0.5 - WLoc.dat.data[0]) < 1e-5
    assert np.abs(0.5 - WLoc.dat.data[1]) < 1e-5


def test_coarsening_localisation_no_localisation():

    mesh = UnitIntervalMesh(1)

    mesh_hierarchy = MeshHierarchy(mesh, 3)

    r_loc = 0

    fs_hierarchy = tuple([FunctionSpace(m, 'DG', 0) for m in mesh_hierarchy])

    WLoc = Function(fs_hierarchy[-1])

    WLoc.dat.data[0] += 1.0

    WLoc_ = Function(fs_hierarchy[-1]).assign(WLoc)

    WLoc = CoarseningLocalisation(WLoc, r_loc)

    assert norm(assemble(WLoc - WLoc_)) == 0


def test_coarsening_localisation_no_hierarchy():

    mesh = UnitIntervalMesh(1)

    fs = FunctionSpace(mesh, 'DG', 0)

    r_locs = [0, 1]

    for r in r_locs:

        WLoc = Function(fs)

        WLoc.dat.data[0] += 1.0

        WLoc_ = Function(fs).assign(WLoc)

        if r == 1:
            with pytest.raises(Exception):
                CoarseningLocalisation(WLoc, r)

        if r == 0:
            WLoc = CoarseningLocalisation(WLoc, r)
            assert norm(assemble(WLoc - WLoc_)) == 0


def test_covariance_localisation_1():

    mesh = UnitIntervalMesh(10)

    vfs = VectorFunctionSpace(mesh, 'DG', 0, dim=10)

    r_loc = 2

    L = CovarianceLocalisation(vfs, r_loc)

    assert np.shape(L.dat.data) == (10, 10)

    assert np.max(np.abs(np.diagonal(L.dat.data) - np.ones(10))) < 1e-5
    assert np.max(np.abs(np.diagonal(L.dat.data, 1) - (0.75 * np.ones(9)))) < 1e-5
    assert np.max(np.abs(np.diagonal(L.dat.data, 2) - (0.25 * np.ones(8)))) < 1e-5
    assert np.max(np.abs(np.diagonal(L.dat.data, -1) - (0.75 * np.ones(9)))) < 1e-5
    assert np.max(np.abs(np.diagonal(L.dat.data, -2) - (0.25 * np.ones(8)))) < 1e-5


def test_covariance_localisation_2():

    mesh = UnitIntervalMesh(10)

    fs = FunctionSpace(mesh, 'CG', 1)
    dim = fs.dof_dset.size
    vfs = VectorFunctionSpace(mesh, 'CG', 1, dim=dim)

    r_loc = 0

    L = CovarianceLocalisation(vfs, r_loc)

    assert np.shape(L.dat.data) == (dim, dim)

    assert np.max(np.abs(np.diagonal(L.dat.data) - np.ones(dim))) < 1e-5


def test_covariance_localisation_3():

    mesh = UnitSquareMesh(10, 10)

    fs = FunctionSpace(mesh, 'DG', 1)
    dim = fs.dof_dset.size
    vfs = VectorFunctionSpace(mesh, 'DG', 1, dim=dim)

    r_loc = 0

    L = CovarianceLocalisation(vfs, r_loc)

    assert np.shape(L.dat.data) == (dim, dim)

    assert np.max(np.abs(np.diagonal(L.dat.data) - np.ones(dim))) < 1e-5


if __name__ == "__main__":
    import os
    pytest.main(os.path.abspath(__file__))
{'fraction_non_alphanumeric': 0.14012534147517275, 'fraction_numerical': 0.033263699180459586, 'mean_word_length': 3.1001317523056655, 'pattern_counts': {'":': 2, '<': 19, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3190678', 'n_tokens_mistral': 2638, 'n_tokens_neox': 2411, 'n_words': 466}