281 lines
9.5 KiB
C
281 lines
9.5 KiB
C
|
|
/*
|
|
* This file is part of the micropython-ulab project,
|
|
*
|
|
* https://github.com/v923z/micropython-ulab
|
|
*
|
|
* The MIT License (MIT)
|
|
*
|
|
* Copyright (c) 2021 Vikas Udupa
|
|
*
|
|
*/
|
|
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include "py/obj.h"
|
|
#include "py/runtime.h"
|
|
#include "py/misc.h"
|
|
|
|
#include "../../ulab.h"
|
|
#include "../../ulab_tools.h"
|
|
#include "../../numpy/linalg/linalg_tools.h"
|
|
#include "linalg.h"
|
|
|
|
#if ULAB_SCIPY_HAS_LINALG_MODULE
|
|
//|
|
|
//| import ulab.scipy
|
|
//| import ulab.numpy
|
|
//|
|
|
//| """Linear algebra functions"""
|
|
//|
|
|
|
|
#if ULAB_MAX_DIMS > 1
|
|
|
|
//| def solve_triangular(A: ulab.numpy.ndarray, b: ulab.numpy.ndarray, lower: bool) -> ulab.numpy.ndarray:
|
|
//| """
|
|
//| :param ~ulab.numpy.ndarray A: a matrix
|
|
//| :param ~ulab.numpy.ndarray b: a vector
|
|
//| :param ~bool lower: if true, use only data contained in lower triangle of A, else use upper triangle of A
|
|
//| :return: solution to the system A x = b. Shape of return matches b
|
|
//| :raises TypeError: if A and b are not of type ndarray and are not dense
|
|
//| :raises ValueError: if A is a singular matrix
|
|
//|
|
|
//| Solve the equation A x = b for x, assuming A is a triangular matrix"""
|
|
//| ...
|
|
//|
|
|
|
|
static mp_obj_t solve_triangular(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
|
|
|
|
size_t i, j;
|
|
|
|
static const mp_arg_t allowed_args[] = {
|
|
{ MP_QSTR_, MP_ARG_REQUIRED | MP_ARG_OBJ, { .u_rom_obj = MP_ROM_NONE} } ,
|
|
{ MP_QSTR_, MP_ARG_REQUIRED | MP_ARG_OBJ, { .u_rom_obj = MP_ROM_NONE} } ,
|
|
{ MP_QSTR_lower, MP_ARG_OBJ, { .u_rom_obj = MP_ROM_TRUE } },
|
|
};
|
|
|
|
mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
|
|
mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
|
|
|
|
if(!mp_obj_is_type(args[0].u_obj, &ulab_ndarray_type) || !mp_obj_is_type(args[1].u_obj, &ulab_ndarray_type)) {
|
|
mp_raise_TypeError(MP_ERROR_TEXT("first two arguments must be ndarrays"));
|
|
}
|
|
|
|
ndarray_obj_t *A = MP_OBJ_TO_PTR(args[0].u_obj);
|
|
ndarray_obj_t *b = MP_OBJ_TO_PTR(args[1].u_obj);
|
|
|
|
if(!ndarray_is_dense(A) || !ndarray_is_dense(b)) {
|
|
mp_raise_TypeError(MP_ERROR_TEXT("input must be a dense ndarray"));
|
|
}
|
|
|
|
size_t A_rows = A->shape[ULAB_MAX_DIMS - 2];
|
|
size_t A_cols = A->shape[ULAB_MAX_DIMS - 1];
|
|
|
|
uint8_t *A_arr = (uint8_t *)A->array;
|
|
uint8_t *b_arr = (uint8_t *)b->array;
|
|
|
|
mp_float_t (*get_A_ele)(void *) = ndarray_get_float_function(A->dtype);
|
|
mp_float_t (*get_b_ele)(void *) = ndarray_get_float_function(b->dtype);
|
|
|
|
uint8_t *temp_A = A_arr;
|
|
|
|
// check if input matrix A is singular
|
|
for (i = 0; i < A_rows; i++) {
|
|
if (MICROPY_FLOAT_C_FUN(fabs)(get_A_ele(A_arr)) < LINALG_EPSILON)
|
|
mp_raise_ValueError(MP_ERROR_TEXT("input matrix is singular"));
|
|
A_arr += A->strides[ULAB_MAX_DIMS - 2];
|
|
A_arr += A->strides[ULAB_MAX_DIMS - 1];
|
|
}
|
|
|
|
A_arr = temp_A;
|
|
|
|
ndarray_obj_t *x = ndarray_new_dense_ndarray(b->ndim, b->shape, NDARRAY_FLOAT);
|
|
mp_float_t *x_arr = (mp_float_t *)x->array;
|
|
|
|
if (mp_obj_is_true(args[2].u_obj)) {
|
|
// Solve the lower triangular matrix by iterating each row of A.
|
|
// Start by finding the first unknown using the first row.
|
|
// On finding this unknown, find the second unknown using the second row.
|
|
// Continue the same till the last unknown is found using the last row.
|
|
|
|
for (i = 0; i < A_rows; i++) {
|
|
mp_float_t sum = 0.0;
|
|
for (j = 0; j < i; j++) {
|
|
sum += (get_A_ele(A_arr) * (*x_arr++));
|
|
A_arr += A->strides[ULAB_MAX_DIMS - 1];
|
|
}
|
|
|
|
sum = (get_b_ele(b_arr) - sum) / (get_A_ele(A_arr));
|
|
*x_arr = sum;
|
|
|
|
x_arr -= j;
|
|
A_arr -= A->strides[ULAB_MAX_DIMS - 1] * j;
|
|
A_arr += A->strides[ULAB_MAX_DIMS - 2];
|
|
b_arr += b->strides[ULAB_MAX_DIMS - 1];
|
|
}
|
|
} else {
|
|
// Solve the upper triangular matrix by iterating each row of A.
|
|
// Start by finding the last unknown using the last row.
|
|
// On finding this unknown, find the last-but-one unknown using the last-but-one row.
|
|
// Continue the same till the first unknown is found using the first row.
|
|
|
|
A_arr += (A->strides[ULAB_MAX_DIMS - 2] * A_rows);
|
|
b_arr += (b->strides[ULAB_MAX_DIMS - 1] * A_cols);
|
|
x_arr += A_cols;
|
|
|
|
for (i = A_rows - 1; i < A_rows; i--) {
|
|
mp_float_t sum = 0.0;
|
|
for (j = i + 1; j < A_cols; j++) {
|
|
sum += (get_A_ele(A_arr) * (*x_arr++));
|
|
A_arr += A->strides[ULAB_MAX_DIMS - 1];
|
|
}
|
|
|
|
x_arr -= (j - i);
|
|
A_arr -= (A->strides[ULAB_MAX_DIMS - 1] * (j - i));
|
|
b_arr -= b->strides[ULAB_MAX_DIMS - 1];
|
|
|
|
sum = (get_b_ele(b_arr) - sum) / get_A_ele(A_arr);
|
|
*x_arr = sum;
|
|
|
|
A_arr -= A->strides[ULAB_MAX_DIMS - 2];
|
|
}
|
|
}
|
|
|
|
return MP_OBJ_FROM_PTR(x);
|
|
}
|
|
|
|
MP_DEFINE_CONST_FUN_OBJ_KW(linalg_solve_triangular_obj, 2, solve_triangular);
|
|
|
|
//| def cho_solve(L: ulab.numpy.ndarray, b: ulab.numpy.ndarray) -> ulab.numpy.ndarray:
|
|
//| """
|
|
//| :param ~ulab.numpy.ndarray L: the lower triangular, Cholesky factorization of A
|
|
//| :param ~ulab.numpy.ndarray b: right-hand-side vector b
|
|
//| :return: solution to the system A x = b. Shape of return matches b
|
|
//| :raises TypeError: if L and b are not of type ndarray and are not dense
|
|
//|
|
|
//| Solve the linear equations A x = b, given the Cholesky factorization of A as input"""
|
|
//| ...
|
|
//|
|
|
|
|
static mp_obj_t cho_solve(mp_obj_t _L, mp_obj_t _b) {
|
|
|
|
if(!mp_obj_is_type(_L, &ulab_ndarray_type) || !mp_obj_is_type(_b, &ulab_ndarray_type)) {
|
|
mp_raise_TypeError(MP_ERROR_TEXT("first two arguments must be ndarrays"));
|
|
}
|
|
|
|
ndarray_obj_t *L = MP_OBJ_TO_PTR(_L);
|
|
ndarray_obj_t *b = MP_OBJ_TO_PTR(_b);
|
|
|
|
if(!ndarray_is_dense(L) || !ndarray_is_dense(b)) {
|
|
mp_raise_TypeError(MP_ERROR_TEXT("input must be a dense ndarray"));
|
|
}
|
|
|
|
mp_float_t (*get_L_ele)(void *) = ndarray_get_float_function(L->dtype);
|
|
mp_float_t (*get_b_ele)(void *) = ndarray_get_float_function(b->dtype);
|
|
void (*set_L_ele)(void *, mp_float_t) = ndarray_set_float_function(L->dtype);
|
|
|
|
size_t L_rows = L->shape[ULAB_MAX_DIMS - 2];
|
|
size_t L_cols = L->shape[ULAB_MAX_DIMS - 1];
|
|
|
|
// Obtain transpose of the input matrix L in L_t
|
|
size_t L_t_shape[ULAB_MAX_DIMS];
|
|
size_t L_t_rows = L_t_shape[ULAB_MAX_DIMS - 2] = L_cols;
|
|
size_t L_t_cols = L_t_shape[ULAB_MAX_DIMS - 1] = L_rows;
|
|
ndarray_obj_t *L_t = ndarray_new_dense_ndarray(L->ndim, L_t_shape, L->dtype);
|
|
|
|
uint8_t *L_arr = (uint8_t *)L->array;
|
|
uint8_t *L_t_arr = (uint8_t *)L_t->array;
|
|
uint8_t *b_arr = (uint8_t *)b->array;
|
|
|
|
size_t i, j;
|
|
|
|
uint8_t *L_ptr = L_arr;
|
|
uint8_t *L_t_ptr = L_t_arr;
|
|
for (i = 0; i < L_rows; i++) {
|
|
for (j = 0; j < L_cols; j++) {
|
|
set_L_ele(L_t_ptr, get_L_ele(L_ptr));
|
|
L_t_ptr += L_t->strides[ULAB_MAX_DIMS - 2];
|
|
L_ptr += L->strides[ULAB_MAX_DIMS - 1];
|
|
}
|
|
|
|
L_t_ptr -= j * L_t->strides[ULAB_MAX_DIMS - 2];
|
|
L_t_ptr += L_t->strides[ULAB_MAX_DIMS - 1];
|
|
L_ptr -= j * L->strides[ULAB_MAX_DIMS - 1];
|
|
L_ptr += L->strides[ULAB_MAX_DIMS - 2];
|
|
}
|
|
|
|
ndarray_obj_t *x = ndarray_new_dense_ndarray(b->ndim, b->shape, NDARRAY_FLOAT);
|
|
mp_float_t *x_arr = (mp_float_t *)x->array;
|
|
|
|
ndarray_obj_t *y = ndarray_new_dense_ndarray(b->ndim, b->shape, NDARRAY_FLOAT);
|
|
mp_float_t *y_arr = (mp_float_t *)y->array;
|
|
|
|
// solve L y = b to obtain y, where L_t x = y
|
|
for (i = 0; i < L_rows; i++) {
|
|
mp_float_t sum = 0.0;
|
|
for (j = 0; j < i; j++) {
|
|
sum += (get_L_ele(L_arr) * (*y_arr++));
|
|
L_arr += L->strides[ULAB_MAX_DIMS - 1];
|
|
}
|
|
|
|
sum = (get_b_ele(b_arr) - sum) / (get_L_ele(L_arr));
|
|
*y_arr = sum;
|
|
|
|
y_arr -= j;
|
|
L_arr -= L->strides[ULAB_MAX_DIMS - 1] * j;
|
|
L_arr += L->strides[ULAB_MAX_DIMS - 2];
|
|
b_arr += b->strides[ULAB_MAX_DIMS - 1];
|
|
}
|
|
|
|
// using y, solve L_t x = y to obtain x
|
|
L_t_arr += (L_t->strides[ULAB_MAX_DIMS - 2] * L_t_rows);
|
|
y_arr += L_t_cols;
|
|
x_arr += L_t_cols;
|
|
|
|
for (i = L_t_rows - 1; i < L_t_rows; i--) {
|
|
mp_float_t sum = 0.0;
|
|
for (j = i + 1; j < L_t_cols; j++) {
|
|
sum += (get_L_ele(L_t_arr) * (*x_arr++));
|
|
L_t_arr += L_t->strides[ULAB_MAX_DIMS - 1];
|
|
}
|
|
|
|
x_arr -= (j - i);
|
|
L_t_arr -= (L_t->strides[ULAB_MAX_DIMS - 1] * (j - i));
|
|
y_arr--;
|
|
|
|
sum = ((*y_arr) - sum) / get_L_ele(L_t_arr);
|
|
*x_arr = sum;
|
|
|
|
L_t_arr -= L_t->strides[ULAB_MAX_DIMS - 2];
|
|
}
|
|
|
|
return MP_OBJ_FROM_PTR(x);
|
|
}
|
|
|
|
MP_DEFINE_CONST_FUN_OBJ_2(linalg_cho_solve_obj, cho_solve);
|
|
|
|
#endif
|
|
|
|
static const mp_rom_map_elem_t ulab_scipy_linalg_globals_table[] = {
|
|
{ MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_linalg) },
|
|
#if ULAB_MAX_DIMS > 1
|
|
#if ULAB_SCIPY_LINALG_HAS_SOLVE_TRIANGULAR
|
|
{ MP_ROM_QSTR(MP_QSTR_solve_triangular), MP_ROM_PTR(&linalg_solve_triangular_obj) },
|
|
#endif
|
|
#if ULAB_SCIPY_LINALG_HAS_CHO_SOLVE
|
|
{ MP_ROM_QSTR(MP_QSTR_cho_solve), MP_ROM_PTR(&linalg_cho_solve_obj) },
|
|
#endif
|
|
#endif
|
|
};
|
|
|
|
static MP_DEFINE_CONST_DICT(mp_module_ulab_scipy_linalg_globals, ulab_scipy_linalg_globals_table);
|
|
|
|
const mp_obj_module_t ulab_scipy_linalg_module = {
|
|
.base = { &mp_type_module },
|
|
.globals = (mp_obj_dict_t*)&mp_module_ulab_scipy_linalg_globals,
|
|
};
|
|
#if CIRCUITPY_ULAB
|
|
MP_REGISTER_MODULE(MP_QSTR_ulab_dot_scipy_dot_linalg, ulab_scipy_linalg_module);
|
|
#endif
|
|
#endif
|