File field.hpp

File List > include > multigrid > field.hpp

Go to the documentation of this file

#pragma once

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <stdexcept>
#include <string>
#include <vector>

#include "multigrid/types.hpp"

namespace mgrid {

class Field2D {
  public:
    Field2D() = default;

    Field2D(int nx, int nz) : nx_{nx}, nz_{nz}, data_(static_cast<std::size_t>(nx) * nz, 0.0) {}

    void resize(int nx, int nz) {
        nx_ = nx;
        nz_ = nz;
        data_.assign(static_cast<std::size_t>(nx) * nz, 0.0);
    }

    [[nodiscard]] int rows() const noexcept { return nx_; }
    [[nodiscard]] int cols() const noexcept { return nz_; }
    [[nodiscard]] int size() const noexcept { return nx_ * nz_; }
    [[nodiscard]] std::array<int, 2> shape() const noexcept { return {nx_, nz_}; }

    [[nodiscard]] double &operator()(int i, int j) { return data_[index(i, j)]; }
    [[nodiscard]] double operator()(int i, int j) const { return data_[index(i, j)]; }
    [[nodiscard]] double &operator[](int i, int j) { return (*this)(i, j); }
    [[nodiscard]] double operator[](int i, int j) const { return (*this)(i, j); }

    [[nodiscard]] double *data() noexcept { return data_.data(); }
    [[nodiscard]] const double *data() const noexcept { return data_.data(); }

    void fill(double value) { std::ranges::fill(data_, value); }

    Field2D &operator+=(double s) {
        for (auto &v : data_) {
            v += s;
        }
        return *this;
    }
    Field2D &operator-=(double s) {
        for (auto &v : data_) {
            v -= s;
        }
        return *this;
    }
    Field2D &operator*=(double s) {
        for (auto &v : data_) {
            v *= s;
        }
        return *this;
    }
    Field2D &operator/=(double s) {
        for (auto &v : data_) {
            v /= s;
        }
        return *this;
    }

    Field2D &operator+=(const Field2D &o) {
        return combine(o, [](double a, double b) { return a + b; });
    }
    Field2D &operator-=(const Field2D &o) {
        return combine(o, [](double a, double b) { return a - b; });
    }
    Field2D &operator*=(const Field2D &o) {
        return combine(o, [](double a, double b) { return a * b; });
    }
    Field2D &operator/=(const Field2D &o) {
        return combine(o, [](double a, double b) { return a / b; });
    }

    [[nodiscard]] Field2D copy() const { return *this; }

    [[nodiscard]] const std::vector<double> &raw() const noexcept { return data_; }
    [[nodiscard]] std::vector<double> &raw() noexcept { return data_; }

  private:
    [[nodiscard]] std::size_t index(int i, int j) const {
        return static_cast<std::size_t>(i) * nz_ + j;
    }

    template <class Op> Field2D &combine(const Field2D &o, Op op) {
        if (o.nx_ != nx_ || o.nz_ != nz_) {
            throw std::invalid_argument("Field2D shape mismatch");
        }
        for (std::size_t k = 0; k < data_.size(); ++k) {
            data_[k] = op(data_[k], o.data_[k]);
        }
        return *this;
    }

    int nx_{0};
    int nz_{0};
    std::vector<double> data_;
};

[[nodiscard]] inline double sum_of_squares(const Field2D &f) {
    double s = 0.0;
    for (double v : f.raw()) {
        s += v * v;
    }
    return s;
}

[[nodiscard]] inline double frobenius_norm(const Field2D &f) {
    return std::sqrt(sum_of_squares(f) / static_cast<double>(f.size()));
}

template <class F> void apply(Field2D &f, F fn) {
    for (double &v : f.raw()) {
        v = fn(v);
    }
}

template <class F> void for_each_index(const Field2D &f, F fn) {
    for (int i = 0; i < f.rows(); ++i) {
        for (int j = 0; j < f.cols(); ++j) {
            fn(i, j);
        }
    }
}

template <class F> void red_black_sweep(const Field2D &f, F fn) {
    for (int colour = 0; colour < 2; ++colour) {
        for (int i = 1; i < f.rows() - 1; ++i) {
            for (int j = 1; j < f.cols() - 1; ++j) {
                if ((i + j) % 2 == colour) {
                    fn(i, j);
                }
            }
        }
    }
}

template <class F>
void red_black_sweep([[maybe_unused]] const Field2D &f, F fn, const Axis &ax, const Axis &az) {
    for (int colour = 0; colour < 2; ++colour) {
        for (int i = ax.lo(); i < ax.hi(); ++i) {
            for (int j = az.lo(); j < az.hi(); ++j) {
                if ((i + j) % 2 == colour) {
                    fn(i, j);
                }
            }
        }
    }
}

inline void project_out_constant(Field2D &f) {
    if (f.size() == 0) {
        return;
    }
    double sum = 0.0;
    for (const double v : f.raw()) {
        sum += v;
    }
    const double mean = sum / static_cast<double>(f.size());
    for (double &v : f.raw()) {
        v -= mean;
    }
}

} // namespace mgrid