Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Small Matrix (6x6, F) #399

Merged
merged 3 commits into from
Jan 6, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions docs/source/usage/api.rst
Original file line number Diff line number Diff line change
@@ -174,6 +174,13 @@ Data Containers
:members:
:undoc-members:

Small Matrices and Vectors
""""""""""""""""""""""""""

.. autoclass:: amrex.space3d.SmallMatrix_6x6_F_SI1_double
:members:
:undoc-members:

Utility
"""""""

1 change: 1 addition & 0 deletions src/Base/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -25,6 +25,7 @@ foreach(D IN LISTS AMReX_SPACEDIM)
IndexType.cpp
IntVect.cpp
RealVect.cpp
SmallMatrix.cpp
MultiFab.cpp
ParallelDescriptor.cpp
ParmParse.cpp
367 changes: 367 additions & 0 deletions src/Base/SmallMatrix.H
Original file line number Diff line number Diff line change
@@ -0,0 +1,367 @@
/* Copyright 2021-2025 The AMReX Community
*
* Authors: Axel Huebl
* License: BSD-3-Clause-LBNL
*/
#pragma once

#include "pyAMReX.H"

#include <AMReX_SmallMatrix.H>

#include <complex>
#include <cstdint>
#include <iterator>
#include <sstream>
#include <type_traits>
#include <vector>


namespace
{
// helper type traits
template <typename T>
struct get_value_type { using value_type = T; };
template <typename T>
struct get_value_type<std::complex<T>> { using value_type = T; };
template <typename T>
using get_value_type_t = typename get_value_type<T>::value_type;

// helper to check if Array4<T> is of constant value type T
template <typename T>
constexpr bool is_not_const ()
{
return std::is_same_v<
std::remove_cv_t<
T
>,
T
> &&
std::is_same_v<
std::remove_cv_t<
get_value_type_t<T>
>,
get_value_type_t<T>
>;
}

/** CPU: __array_interface__ v3
*
* https://numpy.org/doc/stable/reference/arrays.interface.html
*/
template<
class T,
int NRows,
int NCols,
amrex::Order ORDER = amrex::Order::F,
int StartIndex = 0
>
py::dict
array_interface (amrex::SmallMatrix<T, NRows, NCols, ORDER, StartIndex> const & m)
{
using namespace amrex;

auto d = py::dict();
// provide C index order for shape and strides
auto shape = m.ordering == Order::F ? py::make_tuple(
py::ssize_t(NRows),
py::ssize_t(NCols) // fastest varying index
) : py::make_tuple(
py::ssize_t(NCols),
py::ssize_t(NRows) // fastest varying index
);
// buffer protocol strides are in bytes
auto const strides = m.ordering == Order::F ? py::make_tuple(
py::ssize_t(sizeof(T) * NCols),
py::ssize_t(sizeof(T)) // fastest varying index
) : py::make_tuple(
py::ssize_t(sizeof(T) * NRows),
py::ssize_t(sizeof(T)) // fastest varying index
);
bool const read_only = false; // note: we could decide on is_not_const,
// but many libs, e.g. PyTorch, do not
// support read-only and will raise
// warnings, casting to read-write
d["data"] = py::make_tuple(std::intptr_t(&m.template get<0>()), read_only);
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

// note: if we want to keep the same global indexing with non-zero
// box small_end as in AMReX, then we can explore playing with
// this offset as well
//d["offset"] = 0; // default
//d["mask"] = py::none(); // default

d["shape"] = shape;
// we could also set this after checking the strides are C-style contiguous:
//if (is_contiguous<T>(shape, strides))
// d["strides"] = py::none(); // C-style contiguous
//else
d["strides"] = strides;

// type description
// for more complicated types, e.g., tuples/structs
//d["descr"] = ...;
// we currently only need this
using T_no_cv = std::remove_cv_t<T>;
d["typestr"] = py::format_descriptor<T_no_cv>::format();

d["version"] = 3;
return d;
}

template<class SM>
py::class_<SM>
make_SmallMatrix_or_Vector (py::module &m, std::string typestr)
{
using namespace amrex;

using T = typename SM::value_type;
using T_no_cv = std::remove_cv_t<T>;
static constexpr int row_size = SM::row_size;
static constexpr int column_size = SM::column_size;
static constexpr Order ordering = SM::ordering;
static constexpr int starting_index = SM::starting_index;

// dispatch simpler via: py::format_descriptor<T>::format() naming
// but note the _const suffix that might be needed
auto const sm_name = std::string("SmallMatrix_")
.append(std::to_string(row_size)).append("x").append(std::to_string(column_size))
.append("_").append(ordering == Order::F ? "F" : "C")
.append("_SI").append(std::to_string(starting_index))
.append("_").append(typestr);
py::class_< SM > py_sm(m, sm_name.c_str());
py_sm
.def("__repr__",
[sm_name](SM const &) {
return "<amrex." + sm_name + ">";
}
)
.def("__str__",
[sm_name](SM const & sm) {
std::stringstream ss;
ss << sm;
return ss.str();
}
)

.def_property_readonly("size", [](SM const &){ return SM::row_size * SM::column_size; })
.def_property_readonly("row_size", [](SM const &){ return SM::row_size; })
.def_property_readonly("column_size", [](SM const &){ return SM::column_size; })
.def_property_readonly("order", [](SM const &){ return SM::ordering == Order::F ? "F" : "C"; }) // NumPy name
.def_property_readonly("starting_index", [](SM const &){ return SM::starting_index; })

.def_static("zero", [](){ return SM::Zero(); })

.def(py::init([](){ return SM{}; })) // zero-init
.def(py::init<SM const &>()) // copy-init

/* init from a numpy or other buffer protocol array: copy
*/
.def(py::init([](py::array_t<T> & arr)
{
py::buffer_info buf = arr.request();

constexpr bool is_vector = SM::column_size == 1 || SM::row_size == 1;
constexpr int sm_dim = is_vector ? 1 : 2;
if (buf.ndim != sm_dim)
throw std::runtime_error("The SmallMatrix to create is " + std::to_string(sm_dim) +
"D, but the passed array is " + std::to_string(buf.ndim) + "D.");
if (buf.size != SM::column_size * SM::row_size)
throw std::runtime_error("Array size mismatch: Expected " + std::to_string(SM::column_size * SM::row_size) +
" elements, but passed " + std::to_string(buf.size) + " elements.");

if (buf.format != py::format_descriptor<T_no_cv>::format())
throw std::runtime_error("Incompatible format: expected '" +
py::format_descriptor<T_no_cv>::format() +
"' and received '" + buf.format + "'!");

// TODO: check that strides are either exact or None in buf (e.g., F or C contiguous)
// TODO: transpose if SM order is not C?

auto sm = std::make_unique< SM >();
auto * src = static_cast<T*>(buf.ptr);
std::copy(src, src + buf.size, &sm->template get<0>());
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@WeiqunZhang same as above (re: start index not honored in .get)


// todo: we could check and store here if the array buffer we got is read-only

return sm;
}))

/* init from __cuda_array_interface__: device-to-host copy
* TODO
*/


// CPU: __array_interface__ v3
// https://numpy.org/doc/stable/reference/arrays.interface.html
.def_property_readonly("__array_interface__", [](SM const & sm) {
return array_interface(sm);
})

// CPU: __array_function__ interface (TODO)
//
// NEP 18 — A dispatch mechanism for NumPy's high level array functions.
// https://numpy.org/neps/nep-0018-array-function-protocol.html
// This enables code using NumPy to be directly operated on Array4 arrays.
// __array_function__ feature requires NumPy 1.16 or later.


// Nvidia GPUs: __cuda_array_interface__ v3
// https://numba.readthedocs.io/en/latest/cuda/cuda_array_interface.html
.def_property_readonly("__cuda_array_interface__", [](SM const & sm)
{
auto d = array_interface(sm);

// data:
// Because the user of the interface may or may not be in the same context, the most common case is to use cuPointerGetAttribute with CU_POINTER_ATTRIBUTE_DEVICE_POINTER in the CUDA driver API (or the equivalent CUDA Runtime API) to retrieve a device pointer that is usable in the currently active context.
// TODO For zero-size arrays, use 0 here.

// None or integer
// An optional stream upon which synchronization must take place at the point of consumption, either by synchronizing on the stream or enqueuing operations on the data on the given stream. Integer values in this entry are as follows:
// 0: This is disallowed as it would be ambiguous between None and the default stream, and also between the legacy and per-thread default streams. Any use case where 0 might be given should either use None, 1, or 2 instead for clarity.
// 1: The legacy default stream.
// 2: The per-thread default stream.
// Any other integer: a cudaStream_t represented as a Python integer.
// When None, no synchronization is required.
d["stream"] = py::none();

d["version"] = 3;
return d;
})


// TODO: __dlpack__ __dlpack_device__
// DLPack protocol (CPU, NVIDIA GPU, AMD GPU, Intel GPU, etc.)
// https://dmlc.github.io/dlpack/latest/
// https://data-apis.org/array-api/latest/design_topics/data_interchange.html
// https://github.com/data-apis/consortium-feedback/issues/1
// https://github.com/dmlc/dlpack/blob/master/include/dlpack/dlpack.h
// https://docs.cupy.dev/en/stable/user_guide/interoperability.html#dlpack-data-exchange-protocol

;

return py_sm;
}

template<class SM>
void add_matrix_methods (py::class_<SM> & py_sm)
{
using T = typename SM::value_type;
using T_no_cv = std::remove_cv_t<T>;
static constexpr int row_size = SM::row_size;
static constexpr int column_size = SM::column_size;
static constexpr int starting_index = SM::starting_index;

py_sm
.def("dot", &SM::dot)
.def("prod", &SM::product) // NumPy name
.def("set_val", &SM::setVal)
.def("sum", &SM::sum)
.def_property_readonly("T", &SM::transpose) // NumPy name

// operators
.def(py::self + py::self)
.def(py::self - py::self)
.def(py::self * amrex::Real())
.def(amrex::Real() * py::self)
.def(-py::self)

// getter
.def("__getitem__", [](SM & sm, std::array<int, 2> const & key){
if (key[0] < starting_index || key[0] >= row_size + starting_index ||
key[1] < starting_index || key[1] >= column_size + starting_index)
throw std::runtime_error(
"Index out of bounds: [" +
std::to_string(key[0]) + ", " +
std::to_string(key[1]) + "]");
return sm(key[0], key[1]);
})
;

// setter
if constexpr (is_not_const<T>())
{
py_sm
.def("__setitem__", [](SM & sm, std::array<int, 2> const & key, T_no_cv const value){
if (key[0] < SM::starting_index || key[0] >= SM::row_size + SM::starting_index ||
key[1] < SM::starting_index || key[1] >= SM::column_size + SM::starting_index)
{
throw std::runtime_error(
"Index out of bounds: [" +
std::to_string(key[0]) + ", " +
std::to_string(key[1]) + "]");
}
sm(key[0], key[1]) = value;
})
;
}

// square matrix
if constexpr (row_size == column_size)
{
py_sm
.def_static("identity", []() { return SM::Identity(); })
.def("trace", [](SM & sm){ return sm.trace(); })
.def("transpose_in_place", [](SM & sm){ return sm.transposeInPlace(); })
;
}
}

template<class T_SV>
void add_get_set_Vector (py::class_<T_SV> &py_v)
{
using self = T_SV;
using T = typename T_SV::value_type;
using T_no_cv = std::remove_cv_t<T>;

py_v
.def("__getitem__", [](self & sm, int key){
if (key < self::starting_index || key >= self::column_size * self::row_size + self::starting_index)
throw std::runtime_error("Index out of bounds: " + std::to_string(key));
return sm(key);
})
.def("__setitem__", [](self & sm, int key, T_no_cv const value){
if (key < self::starting_index || key >= self::column_size * self::row_size + self::starting_index)
throw std::runtime_error("Index out of bounds: " + std::to_string(key));
sm(key) = value;
})
;
}
}

namespace pyAMReX
{
template<
class T,
int NRows,
int NCols,
amrex::Order ORDER = amrex::Order::F,
int StartIndex = 0
>
void make_SmallMatrix (py::module &m, std::string typestr)
{
using namespace amrex;

using SM = SmallMatrix<T, NRows, NCols, ORDER, StartIndex>;
using SV = SmallMatrix<T, NRows, 1, Order::F, StartIndex>;
using SRV = SmallMatrix<T, 1, NCols, Order::F, StartIndex>;

py::class_<SM> py_sm = make_SmallMatrix_or_Vector<SM>(m, typestr);
py::class_<SV> py_sv = make_SmallMatrix_or_Vector<SV>(m, typestr);
py::class_<SRV> py_srv = make_SmallMatrix_or_Vector<SRV>(m, typestr);

// methods, getter, setter
add_matrix_methods(py_sm);
add_matrix_methods(py_sv);
add_matrix_methods(py_srv);

// vector setter/getter
add_get_set_Vector(py_sv);
add_get_set_Vector(py_srv);

// operators for matrix-matrix & matrix-vector
py_sm
.def(py::self * py::self)
.def(py::self * SV())
.def(SRV() * py::self)
;
}
}
31 changes: 31 additions & 0 deletions src/Base/SmallMatrix.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
/* Copyright 2021-2025 The AMReX Community
*
* Authors: Axel Huebl
* License: BSD-3-Clause-LBNL
*/
#include "pyAMReX.H"

#include "SmallMatrix.H"


void init_SmallMatrix (py::module &m)
{
using namespace pyAMReX;

// 6x6 Matrix as commonly used in accelerator physics
{
constexpr int NRows = 6;
constexpr int NCols = 6;
constexpr amrex::Order ORDER = amrex::Order::F;
constexpr int StartIndex = 1;

make_SmallMatrix< float, NRows, NCols, ORDER, StartIndex >(m, "float");
make_SmallMatrix< double, NRows, NCols, ORDER, StartIndex >(m, "double");
make_SmallMatrix< long double, NRows, NCols, ORDER, StartIndex >(m, "longdouble");
/*
make_SmallMatrix< float const, NRows, NCols, ORDER, StartIndex >(m, "float_const");
make_SmallMatrix< double const, NRows, NCols, ORDER, StartIndex >(m, "double_const");
make_SmallMatrix< long double const, NRows, NCols, ORDER, StartIndex >(m, "longdouble_const");
*/
Comment on lines +25 to +29

Check notice

Code scanning / CodeQL

Commented-out code Note

This comment appears to contain commented-out code.
}
}
142 changes: 142 additions & 0 deletions src/amrex/extensions/SmallMatrix.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,142 @@
"""
This file is part of pyAMReX
Copyright 2025 AMReX community
Authors: Axel Huebl
License: BSD-3-Clause-LBNL
"""


def smallmatrix_to_numpy(self, copy=False, order="F"):
"""
Provide a NumPy view into an SmallMatrix.
Note on the order of indices:
By default, this is as in AMReX in Fortran contiguous order, indexing as
x,y,z. This has performance implications for use in external libraries such
as cupy.
The order="C" option will index as z,y,x and perform better with cupy.
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
Parameters
----------
self : amrex.SmallMatrix_*
A SmallMatrix class in pyAMReX
copy : bool, optional
Copy the data if true, otherwise create a view (default).
order : string, optional
F order (default) or C. C is faster with external libraries.
Returns
-------
np.array
A NumPy 2-dimensional array.
"""
import numpy as np

if copy:
data = np.array(self, copy=True)
else:
data = np.array(self, copy=False)

# TODO: Check self.order == "F" ?
if order == "F":
return data.T
elif order == "C":
return data
else:
raise ValueError("The order argument must be F or C.")


def smallmatrix_to_cupy(self, copy=False, order="F"):
"""
Provide a CuPy view into an SmallMatrix.
Note on the order of indices:
By default, this is as in AMReX in Fortran contiguous order, indexing as
x,y,z. This has performance implications for use in external libraries such
as cupy.
The order="C" option will index as z,y,x and perform better with cupy.
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
Parameters
----------
self : amrex.SmallMatrix_*
A SmallMatrix class in pyAMReX
copy : bool, optional
Copy the data if true, otherwise create a view (default).
order : string, optional
F order (default) or C. C is faster with external libraries.
Returns
-------
cupy.array
A cupy 2-dimensional array.
Raises
------
ImportError
Raises an exception if cupy is not installed
"""
import cupy as cp

# TODO: Check self.order == "F" ?
if order == "F":
return cp.array(self, copy=copy).T
elif order == "C":
return cp.array(self, copy=copy)
else:
raise ValueError("The order argument must be F or C.")


def smallmatrix_to_xp(self, copy=False, order="F"):
"""
Provide a NumPy or CuPy view into a SmallMatrix, depending on amr.Config.have_gpu .
This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
Note on the order of indices:
By default, this is as in AMReX in Fortran contiguous order, indexing as
x,y,z. This has performance implications for use in external libraries such
as cupy.
The order="C" option will index as z,y,x and perform better with cupy.
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
Parameters
----------
self : amrex.SmallMatrix_*
A SmallMatrix class in pyAMReX
copy : bool, optional
Copy the data if true, otherwise create a view (default).
order : string, optional
F order (default) or C. C is faster with external libraries.
Returns
-------
xp.array
A NumPy or CuPy 2-dimensional array.
"""
import inspect

amr = inspect.getmodule(self)
return (
self.to_cupy(copy, order) if amr.Config.have_gpu else self.to_numpy(copy, order)
)


def register_SmallMatrix_extension(amr):
"""SmallMatrix helper methods"""
import inspect
import sys

# register member functions for every Array4_* type
for _, SmallMatrix_type in inspect.getmembers(
sys.modules[amr.__name__],
lambda member: inspect.isclass(member)
and member.__module__ == amr.__name__
and member.__name__.startswith("SmallMatrix_"),
):
SmallMatrix_type.to_numpy = smallmatrix_to_numpy
SmallMatrix_type.to_cupy = smallmatrix_to_cupy
SmallMatrix_type.to_xp = smallmatrix_to_xp
2 changes: 2 additions & 0 deletions src/amrex/space1d/__init__.py
Original file line number Diff line number Diff line change
@@ -50,11 +50,13 @@ def Print(*args, **kwargs):
from ..extensions.MultiFab import register_MultiFab_extension
from ..extensions.ParticleContainer import register_ParticleContainer_extension
from ..extensions.PODVector import register_PODVector_extension
from ..extensions.SmallMatrix import register_SmallMatrix_extension
from ..extensions.StructOfArrays import register_SoA_extension

register_Array4_extension(amrex_1d_pybind)
register_MultiFab_extension(amrex_1d_pybind)
register_PODVector_extension(amrex_1d_pybind)
register_SmallMatrix_extension(amrex_1d_pybind)
register_SoA_extension(amrex_1d_pybind)
register_AoS_extension(amrex_1d_pybind)
register_ParticleContainer_extension(amrex_1d_pybind)
2 changes: 2 additions & 0 deletions src/amrex/space2d/__init__.py
Original file line number Diff line number Diff line change
@@ -50,11 +50,13 @@ def Print(*args, **kwargs):
from ..extensions.MultiFab import register_MultiFab_extension
from ..extensions.ParticleContainer import register_ParticleContainer_extension
from ..extensions.PODVector import register_PODVector_extension
from ..extensions.SmallMatrix import register_SmallMatrix_extension
from ..extensions.StructOfArrays import register_SoA_extension

register_Array4_extension(amrex_2d_pybind)
register_MultiFab_extension(amrex_2d_pybind)
register_PODVector_extension(amrex_2d_pybind)
register_SmallMatrix_extension(amrex_2d_pybind)
register_SoA_extension(amrex_2d_pybind)
register_AoS_extension(amrex_2d_pybind)
register_ParticleContainer_extension(amrex_2d_pybind)
2 changes: 2 additions & 0 deletions src/amrex/space3d/__init__.py
Original file line number Diff line number Diff line change
@@ -50,11 +50,13 @@ def Print(*args, **kwargs):
from ..extensions.MultiFab import register_MultiFab_extension
from ..extensions.ParticleContainer import register_ParticleContainer_extension
from ..extensions.PODVector import register_PODVector_extension
from ..extensions.SmallMatrix import register_SmallMatrix_extension
from ..extensions.StructOfArrays import register_SoA_extension

register_Array4_extension(amrex_3d_pybind)
register_MultiFab_extension(amrex_3d_pybind)
register_PODVector_extension(amrex_3d_pybind)
register_SmallMatrix_extension(amrex_3d_pybind)
register_SoA_extension(amrex_3d_pybind)
register_AoS_extension(amrex_3d_pybind)
register_ParticleContainer_extension(amrex_3d_pybind)
15 changes: 10 additions & 5 deletions src/pyAMReX.cpp
Original file line number Diff line number Diff line change
@@ -14,6 +14,7 @@
// forward declarations of exposed classes
void init_Algorithm(py::module&);
void init_AMReX(py::module&);
void init_AmrMesh(py::module &);
void init_Arena(py::module&);
void init_Array4(py::module&);
void init_BaseFab(py::module&);
@@ -27,22 +28,22 @@ void init_FArrayBox(py::module&);
void init_Geometry(py::module&);
void init_IndexType(py::module &);
void init_IntVect(py::module &);
void init_RealVect(py::module &);
void init_AmrMesh(py::module &);
#ifdef AMREX_USE_MPI
void init_MPMD(py::module &);
#endif
void init_MultiFab(py::module &);
void init_ParallelDescriptor(py::module &);
void init_ParmParse(py::module &);
void init_ParticleContainer(py::module &);
void init_Periodicity(py::module &);
void init_PlotFileUtil(py::module &);
void init_PODVector(py::module &);
void init_RealVect(py::module &);
void init_SmallMatrix(py::module &);
void init_Utility(py::module &);
void init_Vector(py::module &);
void init_Version(py::module &);
void init_VisMF(py::module &);
#ifdef AMREX_USE_MPI
void init_MPMD(py::module &);
#endif

#if AMREX_SPACEDIM == 1
PYBIND11_MODULE(amrex_1d_pybind, m) {
@@ -81,13 +82,15 @@ PYBIND11_MODULE(amrex_3d_pybind, m) {
Periodicity
PlotFileUtil
PODVector
SmallMatrix
StructOfArrays
Utility
Vector
VisMF
)pbdoc";

// note: order from parent to child classes and argument usage

init_AMReX(m);
init_Arena(m);
init_Dim3(m);
@@ -98,6 +101,7 @@ PYBIND11_MODULE(amrex_3d_pybind, m) {
init_Box(m);
init_Periodicity(m);
init_Array4(m);
init_SmallMatrix(m);
init_BoxArray(m);
init_ParmParse(m);
init_CoordSys(m);
@@ -117,6 +121,7 @@ PYBIND11_MODULE(amrex_3d_pybind, m) {
#ifdef AMREX_USE_MPI
init_MPMD(m);
#endif

// Wrappers around standalone functions
init_PlotFileUtil(m);
init_Utility(m);
260 changes: 260 additions & 0 deletions tests/test_smallmatrix.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,260 @@
# -*- coding: utf-8 -*-

import numpy as np
import pytest

import amrex.space3d as amr


def test_smallmatrix():
m66 = amr.SmallMatrix_6x6_F_SI1_double(
[
[1, 2, 3, 4, 5, 6],
[7, 8, 9, 10, 11, 12],
[13, 14, 15, 16, 17, 18],
[19, 20, 21, 22, 23, 24],
[25, 26, 27, 28, 29, 30],
[31, 32, 33, 34, 35, 36],
]
)
v = 1
for j in range(1, 7):
for i in range(1, 7):
assert m66[i, j] == v
v += 1


def test_smallvector():
cv1 = amr.SmallMatrix_6x1_F_SI1_double()
rv1 = amr.SmallMatrix_1x6_F_SI1_double()
cv2 = amr.SmallMatrix_6x1_F_SI1_double([1, 2, 3, 4, 5, 6])
rv2 = amr.SmallMatrix_1x6_F_SI1_double([0, 10, 20, 30, 40, 50])
cv3 = amr.SmallMatrix_6x1_F_SI1_double([0, 1, 2, 3, 4, 5])

for j in range(1, 7):
assert cv1[j] == 0.0
assert rv1[j] == 0.0
assert cv2[j] == j
assert amr.almost_equal(rv2[j], (j - 1) * 10.0)
assert amr.almost_equal(cv3[j], j - 1.0)


def test_smallmatrix_zero():
zero = amr.SmallMatrix_6x6_F_SI1_double()

# Check properties
assert zero.size == 36
assert zero.row_size == 6
assert zero.column_size == 6
assert zero.order == "F"
assert zero.starting_index == 1

# Check values
assert zero.sum() == 0
assert zero.prod() == 0
assert zero.trace() == 0

# assign empty
zeroc = amr.SmallMatrix_6x6_F_SI1_double(zero)

# Check values
assert zeroc.sum() == 0
assert zeroc.prod() == 0
assert zeroc.trace() == 0

# create zero
zerov = amr.SmallMatrix_6x6_F_SI1_double.zero()

# Check values
assert zerov.sum() == 0
assert zerov.prod() == 0
assert zerov.trace() == 0


def test_smallmatrix_identity():
iden = amr.SmallMatrix_6x6_F_SI1_double.identity()

# Check properties
assert iden.size == 36
assert iden.row_size == 6
assert iden.column_size == 6
assert iden.order == "F"
assert iden.starting_index == 1

# Check values
assert iden.sum() == 6
assert iden.prod() == 0
assert iden.trace() == 6


def test_smallmatrix_from_np():
# from numpy (copy)
x = np.ones(
(
6,
6,
)
)
print(f"\nx: {x.__array_interface__} {x.dtype}")
sm = amr.SmallMatrix_6x6_F_SI1_double(x)
print(f"sm: {sm.__array_interface__}")
print(sm)

assert sm.sum() == 36
assert sm.prod() == 1
assert sm.trace() == 6


def test_smallmatrix_to_np():
iden = amr.SmallMatrix_6x6_F_SI1_double.identity()

x = iden.to_numpy()
print(x)

assert x.sum() == 6
assert x.prod() == 0
assert x.trace() == 6
assert not x.flags["C_CONTIGUOUS"]
assert x.flags["F_CONTIGUOUS"]


def test_smallmatrix_smallvector():
v3 = amr.SmallMatrix_6x1_F_SI1_double.zero()
v3[1] = 1.0
v3[2] = 2.0
v3[3] = 3.0
v3[4] = 4.0
v3[5] = 5.0
v3[6] = 6.0
m66 = amr.SmallMatrix_6x6_F_SI1_double.identity()
r = m66 * v3

for i in range(1, 7):
assert amr.almost_equal(r[i], v3[i])


def test_smallmatrix_smallmatrix():
A = amr.SmallMatrix_6x6_F_SI1_double(
[
[1, 0, 1, 0, 1, 0],
[2, 1, 1, 1, 1, 2],
[0, 1, 1, 1, 1, 0],
[1, 1, 2, 2, 1, 1],
[2, 1, 2, 2, 1, 2],
[0, 1, 1, 1, 1, 0],
]
)
B = amr.SmallMatrix_6x6_F_SI1_double(
[
[1, 2, 2, 2, 1, 1],
[2, 3, 1, 1, 1, 3],
[4, 2, 2, 2, 2, 0],
[1, 4, 3, 2, 0, 1],
[2, 3, 1, 0, 0, 2],
[0, 1, 1, 1, 4, 0],
]
)
C = amr.SmallMatrix_6x1_F_SI1_double([10, 8, 6, 4, 2, 0])
ABC = A * B * C
assert ABC[1, 1] == 322
assert ABC[2, 1] == 252
assert ABC[3, 1] == 388
assert ABC[4, 1] == 330
assert ABC[5, 1] == 310
assert ABC[6, 1] == 264

# transpose
CR = amr.SmallMatrix_1x6_F_SI1_double([10, 8, 6, 4, 2, 0])
ABC_T = A.T * B.transpose_in_place() * CR.T
assert ABC_T[1, 1] == 178
assert ABC_T[2, 1] == 402
assert ABC_T[3, 1] == 254
assert ABC_T[4, 1] == 476
assert ABC_T[5, 1] == 550
assert ABC_T[6, 1] == 254


def test_smallmatrix_sum_prod():
m = amr.SmallMatrix_6x6_F_SI1_double()
m.set_val(2.0)

assert m.prod() == 2 ** (m.row_size * m.column_size)
assert m.sum() == 2 * m.row_size * m.column_size


def test_smallmatrix_trace():
m = amr.SmallMatrix_6x6_F_SI1_double(
[
[1.0, 3.4, 4.5, 5.6, 6.7, 7.8],
[1.3, 2.0, 3.4, 4.5, 5.6, 6.7],
[1.3, 1.0, 3.0, 4.5, 5.6, 6.7],
[1.3, 1.4, 4.5, 4.0, 5.6, 6.7],
[1.3, 1.0, 4.5, 5.6, 5.0, 6.7],
[1.3, 1.4, 3.0, 4.5, 6.7, 6.0],
]
)
assert m.trace() == 1.0 + 2.0 + 3.0 + 4.0 + 5.0 + 6.0


def test_smallmatrix_scalar():
A = amr.SmallMatrix_6x6_F_SI1_double(
[
[+1.0, +2, +3, +4, +5, +6],
[+7, +8, +9, +10, +11, +12],
[+13, +14, +15, +16, +17, +18],
[+19, +20, +21, +22, +23, +24],
[+25, +26, +27, +28, +29, +30],
[+31, +32, +33, +34, +35, +36],
]
)
B = amr.SmallMatrix_6x6_F_SI1_double(A)
B *= -1.0

# test matrix-scalar and scalar-matrix
C = A * 2.0 + 2.0 * B
assert np.allclose(C.to_numpy(), 0.0)

# test unary- operator and point-wise minus
D = -A - B
assert np.allclose(D.to_numpy(), 0.0)

# dot product
E = amr.SmallMatrix_6x6_F_SI1_double()
E.set_val(-1.0)
assert A.dot(E) == -666


def test_smallmatrix_rangecheck():
cv = amr.SmallMatrix_6x1_F_SI1_double()
rv = amr.SmallMatrix_1x6_F_SI1_double()
m66 = amr.SmallMatrix_6x6_F_SI1_double(
[
[1, 2, 3, 4, 5, 6],
[7, 8, 9, 10, 11, 12],
[13, 14, 15, 16, 17, 18],
[19, 20, 21, 22, 23, 24],
[25, 26, 27, 28, 29, 30],
[31, 32, 33, 34, 35, 36],
]
)

with pytest.raises(RuntimeError):
cv[0]
with pytest.raises(RuntimeError):
cv[7]
with pytest.raises(RuntimeError):
rv[0]
with pytest.raises(RuntimeError):
rv[7]
with pytest.raises(RuntimeError):
m66[0, 0]
with pytest.raises(RuntimeError):
m66[0, 1]
with pytest.raises(RuntimeError):
m66[1, 0]
with pytest.raises(RuntimeError):
m66[7, 7]
with pytest.raises(RuntimeError):
m66[6, 7]
with pytest.raises(RuntimeError):
m66[7, 6]