reference_implementation.py

import numpy as np
import numpy.linalg as la  # noqa: F401
import pyopencl as cl  # noqa: F401
import pyopencl.array  # noqa
import pyopencl.tools  # noqa
import pyopencl.clrandom  # noqa
import loopy as lp  # noqa

import utilities as u
import ideal_gas


def pointwise_fluxes(states):
    return np.array([ideal_gas.flux(s) for s in states.T])


def roe_eigensystem(state_pair, frozen_metrics, direction):
    # FIXME: startup for test suite is pretty slow due to this routine
    #   -- can we speed this up?

    nvars = state_pair.shape[0]
    ndim = frozen_metrics.shape[0]

    prg = u.get_weno_program_with_root_kernel("roe_eigensystem")
    queue = u.get_queue(cl._csc)

    R_dev = u.empty_array_on_device(queue, nvars, nvars)
    R_inv_dev = u.empty_array_on_device(queue, nvars, nvars)
    lam_dev = u.empty_array_on_device(queue, nvars)

    prg(queue, nvars=nvars, ndim=ndim, d=direction,
            states=state_pair, metrics_frozen=frozen_metrics,
            R=R_dev, R_inv=R_inv_dev, lambda_roe=lam_dev)

    return R_dev.get(), R_inv_dev.get(), lam_dev.get()


def lambda_pointwise(states, metrics, direction):
    # FIXME: lose some arguments here
    def lam(state, metric_norm, direction):
        c = ideal_gas.sound_speed(state)
        c_norm = c*metric_norm[direction]
        vel = ideal_gas.velocity(state)[direction]

        result = np.repeat(vel, state.size)
        result[-2] += c_norm
        result[-1] -= c_norm

        return result

    metric_norm = np.sqrt((metrics**2).sum(axis=1))

    return u.transposed_array(
            [lam(s, m, direction-1) for s, m in zip(states.T, metric_norm)])


def wavespeeds(pointwise, roe):
    lam = np.c_[pointwise, roe]
    return 1.1*np.max(np.abs(lam), axis=1)


def split_char_fluxes(states, wavespeeds, frozen_metrics, frozen_jacobian, R_inv):
    def split(flux, state):
        generalized_fluxes = np.dot(flux, frozen_metrics)
        weighted_states = np.outer(wavespeeds, state/frozen_jacobian)

        return (0.5*np.sum(R_inv*(generalized_fluxes + weighted_states), axis=1),
                0.5*np.sum(R_inv*(generalized_fluxes - weighted_states), axis=1))

    fluxes = pointwise_fluxes(states)

    char_fluxes_pos, char_fluxes_neg = zip(
            *[split(f, s) for f, s in zip(fluxes, states.T)])

    return (u.transposed_array(char_fluxes_pos),
            u.transposed_array(char_fluxes_neg))


def oscillation(fluxes):
    sum1 = u.empty_array(fluxes.shape[0], 3)
    sum2 = u.empty_array(fluxes.shape[0], 3)

    sum1[:,0] = fluxes[:,0] - 4*fluxes[:,1] + 3*fluxes[:,2]
    sum1[:,1] = -fluxes[:,1] + fluxes[:,3]
    sum1[:,2] = -3*fluxes[:,2] + 4*fluxes[:,3] - fluxes[:,4]

    for i in range(3):
        sum2[:,i] = fluxes[:,i] - 2*fluxes[:,i+1] + fluxes[:,i+2]

    result = (1.0/4)*(sum1**2) + (13.0/12)*(sum2**2)

    return result


def weno_weights(oscillation, frozen_metric):
    linear = np.array([0.1, 0.6, 0.3])
    eps = 1e-6*frozen_metric

    raw_weights = u.empty_array(5,3)
    for i in range(5):
        for j in range(3):
            raw_weights[i,j] = linear[j]/(oscillation[i,j] + eps)**2

    weight_sum = raw_weights.sum(axis=1)
    weights = u.empty_array(5,3)
    for i in range(5):
        for j in range(3):
            weights[i,j] = raw_weights[i,j]/weight_sum[i]

    return weights


def flux_differences(f):
    w = np.array([-1, 3, -3, 1])
    res = u.empty_array(5, 3)
    for j in range(res.shape[1]):
        res[:,j] = f[:,j:j+4]@w
    return res


def combination_weighting(w):
    return np.array(
            [20*w[:,0] - 1, -10*(w[:,0] + w[:,1]) + 5, np.ones(w.shape[0])]
            ).T


def combine_fluxes(w, f):
    cw = combination_weighting(w)
    return np.sum(cw*f, axis=1)


def dissipation_part(R, char_fluxes, w, sign):
    flux_diff = flux_differences(char_fluxes)[:,::sign]
    flux_comb = combine_fluxes(w, flux_diff)

    return -sign*R@flux_comb/60


def consistent_part(fluxes):
    w = np.array([1.0, -8.0, 37.0, 37.0, -8.0, 1.0])/60.0
    return np.dot(fluxes, w)


def weno_flux(consistent, dissipation_pos, dissipation_neg):
    return consistent + dissipation_pos + dissipation_neg