data_for_test.py

import numpy as np
import pytest
import utilities as u
import weno_reference_implementation as ref


# {{{ FluxDataSingle

class FluxDataSingle:
    # FIXME: can we set some of these constants from ref.gas?
    #   -- if all nvars references come from there, it's relatively easy to
    #   introduce a new gas with more (e.g. scalar) variables
    nvars = 5
    ndim = 3
    dirs = {"x": 1, "y": 2, "z": 3}

    def __init__(self, states_str, direction):

        self.direction = self.dirs[direction]
        self.dir_internal = self.direction-1

        self.metrics = np.array([np.identity(self.ndim) for i in range(6)],
                dtype=np.float64, order="F")
        self.jacobians = np.repeat(1.0, 6)

        # FIXME: should be computed directly from the metrics and jacobians
        self.frozen_metrics = np.mean(self.metrics[2:4], axis=0)
        self.frozen_jacobian = np.mean(self.jacobians[2:4], axis=0)
        self.combined_frozen_metrics = 1.0

        # FIXME: Move array_from_string stuff outside FluxDataSingle
        #   -- just pass an array & have external utilities that generate
        #   Riemann, sine wave, etc. initial conditions
        # FIXME: Consider handling row swapping outside as well?
        # FIXME: Do we even need to swap rows?
        self.state_pair = self.swap_array_rows(
                u.transposed_array_from_string(states_str), self.dir_internal)
        self.states = u.expand_to_n(self.state_pair, 6)

        # FIXME: these should be generalized fluxes
        # FIXME: make a clear distinction between fluxes in physical and
        # generalized coordinates
        self.flux_pair = ref.pointwise_fluxes(
                self.state_pair)[:,:,self.dir_internal].T.copy(order="F")
        self.fluxes = ref.pointwise_fluxes(
                self.states)[:,:,self.dir_internal].T.copy(order="F")

        self.lam_pointwise = ref.lambda_pointwise(
                self.states, self.metrics, self.dir_internal)
        self.R, self.R_inv, self.lam_roe = ref.roe_eigensystem(
                self.state_pair, self.frozen_metrics, self.direction)
        self.wavespeeds = ref.wavespeeds(self.lam_pointwise, self.lam_roe)

        self.char_fluxes_pos, self.char_fluxes_neg = ref.split_char_fluxes(
                self.states, self.wavespeeds,
                self.frozen_metrics[self.dir_internal], self.frozen_jacobian,
                self.R_inv)

        self.oscillation_pos = ref.oscillation(self.char_fluxes_pos)
        self.oscillation_neg = ref.oscillation(self.char_fluxes_neg[:,::-1])

        self.weno_weights_pos = ref.weno_weights(
                self.oscillation_pos, self.combined_frozen_metrics)
        self.weno_weights_neg = ref.weno_weights(
                self.oscillation_neg, self.combined_frozen_metrics)

        self.consistent = ref.consistent_part(self.fluxes)
        self.dissipation_pos = ref.dissipation_part(
                self.R, self.char_fluxes_pos, self.weno_weights_pos, 1)
        self.dissipation_neg = ref.dissipation_part(
                self.R, self.char_fluxes_neg, self.weno_weights_neg, -1)

        self.weno_flux = ref.weno_flux(
                self.consistent, self.dissipation_pos, self.dissipation_neg)

    def swap_array_rows(self, arr, d):
        p = self.permutation(d)
        arr[p, :] = arr[[1, 2, 3], :]
        return arr

    def permutation(self, d):
        return [(d+i) % 3 + 1 for i in range(3)]

# }}}

# {{{ FluxDataVector

# FIXME: is there a better way to divide responsibilities with these fixture classes?
class FluxDataVector:
    nvars = 5
    ndim = 3
    dirs = {"x": 1, "y": 2, "z": 3}
    halo = 3

    def __init__(self, nx, ny, nz, states_str, direction):
        self.direction = self.dirs[direction]
        self.dir_internal = self.direction-1

        self.nx = nx
        self.ny = ny
        self.nz = nz

        self.nxhalo = self.nx + 2*self.halo
        self.nyhalo = self.ny + 2*self.halo
        self.nzhalo = self.nz + 2*self.halo

        self.flux_dims = (self.nvars, self.nx, self.ny, self.nz)

        self.metrics = np.stack(
                [np.stack(
                    [np.stack(
                        [np.identity(self.ndim) for i in range(self.nxhalo)],
                        axis=-1) for j in range(self.nyhalo)],
                    axis=-1) for k in range(self.nzhalo)],
                axis=-1).copy(order="F")
        self.jacobians = np.ones((self.nxhalo, self.nyhalo, self.nzhalo),
                order="F")

        state_pair = self.swap_array_rows(
                u.transposed_array_from_string(states_str), self.direction)
        # FIXME: Move array_from_string stuff outside FluxDataSingle
        #   -- just pass an array & have external utilities that generate
        #   Riemann, sine wave, etc. initial conditions
        # FIXME: Consider handling row swapping outside as well?
        # FIXME: Do we even need to swap rows?
        self.state_pair = self.swap_array_rows(
                u.transposed_array_from_string(states_str), self.dir_internal)
        # NOTE: dimensions are nvars x nxhalo x nyhalo x nzhalo
        self.states = self.fill_from_pair()

        # NOTE: dimensions are nvars x nxhalo x nyhalo x nzhalo
        # FIXME: these should be generalized fluxes
        # FIXME: make a clear distinction between fluxes in physical and
        # generalized coordinates
        npoints = self.nxhalo*self.nyhalo*self.nzhalo
        flat_states = self.states.reshape((self.nvars, npoints))
        self.fluxes = ref.pointwise_fluxes(
                flat_states)[:,:,self.dir_internal].T.reshape(
                        (self.nvars, self.nxhalo, self.nyhalo, self.nzhalo)
                        ).copy(order="F")

        # FIXME: use reference implementation
        # NOTE: dimensions are nvars x nx x ny x nz
        self.flux_derivatives = np.zeros((self.nvars, self.nx, self.ny,
            self.nz), order="F")

    def swap_array_rows(self, arr, d):
        p = self.permutation(d)
        arr[p, :] = arr[[1, 2, 3], :]
        return arr

    def permutation(self, d):
        return [(d-1+i) % 3 + 1 for i in range(3)]

    def fill_from_pair(self):
        d = self.dir_internal
        nx_arr = np.array([self.nxhalo, self.nyhalo, self.nzhalo])
        result = u.expand_to_n(self.state_pair, nx_arr[d])

        for i in range(d):
            result = self.add_dimension(result, nx_arr[i])
            result = np.swapaxes(result, -2, -1)
        for i in range(d+1,self.ndim):
            result = self.add_dimension(result, nx_arr[i])

        return result.copy(order="F")

    def add_dimension(self, arr, n):
        return np.stack([arr for i in range(n)], axis=-1)

# }}}


single_data = {}

single_data["Case flat:x"] = FluxDataSingle(
        states_str="1 1 1 1 5.5,1 1 1 1 5.5", direction="x")
single_data["Case flat:y"] = FluxDataSingle(
        states_str="1 1 1 1 5.5,1 1 1 1 5.5", direction="y")
single_data["Case flat:z"] = FluxDataSingle(
        states_str="1 1 1 1 5.5,1 1 1 1 5.5", direction="z")

single_data["Case a:x"] = FluxDataSingle(
        states_str="2 4 4 4 20,1 1 1 1 5.5", direction="x")
single_data["Case a:y"] = FluxDataSingle(
        states_str="2 4 4 4 20,1 1 1 1 5.5", direction="y")
single_data["Case a:z"] = FluxDataSingle(
        states_str="2 4 4 4 20,1 1 1 1 5.5", direction="z")

single_data["Case b:x"] = FluxDataSingle(
        states_str="1 -1 -1 -1 5.5,2 -4 -4 -4 20", direction="x")
single_data["Case b:y"] = FluxDataSingle(
        states_str="1 -1 -1 -1 5.5,2 -4 -4 -4 20", direction="y")
single_data["Case b:z"] = FluxDataSingle(
        states_str="1 -1 -1 -1 5.5,2 -4 -4 -4 20", direction="z")

single_data["Case c:x"] = FluxDataSingle(
        states_str="2 4 8 12 64,1 1 2 3 11", direction="x")
single_data["Case c:y"] = FluxDataSingle(
        states_str="2 8 12 4 64,1 2 3 1 11", direction="y")
single_data["Case c:z"] = FluxDataSingle(
        states_str="2 12 4 8 64,1 3 1 2 11", direction="z")

single_data["Case d:x"] = FluxDataSingle(
        states_str="1 -1 -2 -3 11,2 -4 -8 -12 64", direction="x")
single_data["Case d:y"] = FluxDataSingle(
        states_str="1 -2 -3 -1 11,2 -8 -12 -4 64", direction="y")
single_data["Case d:z"] = FluxDataSingle(
        states_str="1 -3 -1 -2 11,2 -12 -4 -8 64", direction="z")


vector_data = {}

vector_data["Case flat:x"] = FluxDataVector(
        nx=6, ny=2, nz=2,
        states_str="1 1 1 1 5.5,1 1 1 1 5.5",
        direction="x")
vector_data["Case flat:y"] = FluxDataVector(
        nx=2, ny=6, nz=2,
        states_str="1 1 1 1 5.5,1 1 1 1 5.5",
        direction="y")
vector_data["Case flat:z"] = FluxDataVector(
        nx=2, ny=2, nz=6,
        states_str="1 1 1 1 5.5,1 1 1 1 5.5",
        direction="z")

vector_data["Case a:x"] = FluxDataVector(
        nx=6, ny=2, nz=2,
        states_str="2 4 4 4 20,1 1 1 1 5.5",
        direction="x")
vector_data["Case a:y"] = FluxDataVector(
        nx=2, ny=6, nz=2,
        states_str="2 4 4 4 20,1 1 1 1 5.5",
        direction="y")
vector_data["Case a:z"] = FluxDataVector(
        nx=2, ny=2, nz=6,
        states_str="2 4 4 4 20,1 1 1 1 5.5",
        direction="z")


@pytest.fixture(scope="session", params=[
    "Case flat:x", "Case flat:y", "Case flat:z",
    "Case a:x", "Case a:y", "Case a:z",
    "Case b:x", "Case b:y", "Case b:z",
    "Case c:x", "Case c:y", "Case c:z",
    "Case d:x", "Case d:y", "Case d:z"])
def flux_test_data_fixture(request):
    return single_data[request.param]


@pytest.fixture(scope="session", params=[
    "Case flat:x", "Case flat:y", "Case flat:z"])
def cfd_test_data_fixture(request):
    return vector_data[request.param]


# vim: foldmethod=marker