int luaC_lunum_zeros(lua_State *L)
{
if (lua_isnumber(L, 1)) {
const int N = luaL_checkinteger(L, 1);
const enum ArrayType T = (enum ArrayType) luaL_optinteger(L, 2, ARRAY_TYPE_DOUBLE);
struct Array A = array_new_zeros(N, T);
lunum_pusharray1(L, &A);
return 1;
}
else if (lua_istable(L, 1) || lunum_hasmetatable(L, 1, "array")) {
int Nd;
int *N = (int*) lunum_checkarray2(L, 1, ARRAY_TYPE_INT, &Nd);
const enum ArrayType T = (enum ArrayType) luaL_optinteger(L, 2, ARRAY_TYPE_DOUBLE);
int ntot = 1;
for (int d=0; d<Nd; ++d) ntot *= N[d];
struct Array A = array_new_zeros(ntot, T);
array_resize(&A, N, Nd);
lunum_pusharray1(L, &A);
return 1;
}
else {
luaL_error(L, "argument must be either number, table, or array");
return 0;
}
}
static int luaC_lunum_zeros(lua_State *L)
{
if (lua_isnumber(L, 1)) {
const lua_Integer N = luaL_checkinteger(L, 1);
if (N <= 0) { return luaL_error(L, "Invalid size %d", N); }
const ArrayType T = (ArrayType) luaL_optinteger(L, 2, ARRAY_TYPE_DOUBLE);
Array A = array_new_zeros(N, T);
lunum_pusharray1(L, &A);
return 1;
}
else if (lua_istable(L, 1) || lunum_hasmetatable(L, 1, "array")) {
size_t Nd_t;
size_t *N = (size_t*) lunum_checkarray2(L, 1, ARRAY_TYPE_SIZE_T, &Nd_t);
int Nd = (int)Nd_t;
const ArrayType T = (ArrayType) luaL_optinteger(L, 2, ARRAY_TYPE_DOUBLE);
size_t ntot = 1;
for (int d=0; d<Nd; ++d) ntot *= N[d];
Array A = array_new_zeros(ntot, T);
array_resize_t(&A, N, Nd);
lunum_pusharray1(L, &A);
return 1;
}
else {
return luaL_error(L, "argument must be either number, table, or array");
}
}
void _unary_func(lua_State *L, double(*f)(double), Complex(*g)(Complex), int cast)
{
if (lua_isnumber(L, 1)) {
const double x = lua_tonumber(L, 1);
lua_pushnumber(L, f(x));
}
else if (lunum_hasmetatable(L, 1, "complex")) {
if (g == NULL) {
luaL_error(L, "complex operation not supported");
}
const Complex z = lunum_checkcomplex(L, 1);
lunum_pushcomplex(L, g(z));
}
else if (lunum_hasmetatable(L, 1, "array")) {
struct Array *A = (struct Array*) lunum_checkarray1(L, 1);
if (cast == 0) {
struct Array B = array_new_copy(A, A->dtype);
switch (B.dtype) {
case ARRAY_TYPE_BOOL : EXPR_EVALF(Bool , B.size, B.data); break;
case ARRAY_TYPE_CHAR : EXPR_EVALF(char , B.size, B.data); break;
case ARRAY_TYPE_SHORT : EXPR_EVALF(short , B.size, B.data); break;
case ARRAY_TYPE_INT : EXPR_EVALF(long , B.size, B.data); break;
case ARRAY_TYPE_LONG : EXPR_EVALF(int , B.size, B.data); break;
case ARRAY_TYPE_FLOAT : EXPR_EVALF(float , B.size, B.data); break;
case ARRAY_TYPE_DOUBLE : EXPR_EVALF(double , B.size, B.data); break;
case ARRAY_TYPE_COMPLEX : EXPR_EVALG(Complex, B.size, B.data); break;
}
lunum_pusharray1(L, &B);
}
else if (A->dtype <= ARRAY_TYPE_DOUBLE) {
struct Array B = array_new_copy(A, ARRAY_TYPE_DOUBLE);
double *b = (double*) B.data;
for (int i=0; i<B.size; ++i) b[i] = f(b[i]);
lunum_pusharray1(L, &B);
}
else if (A->dtype == ARRAY_TYPE_COMPLEX) {
if (g == NULL) {
luaL_error(L, "complex operation not supported");
}
struct Array B = array_new_copy(A, ARRAY_TYPE_COMPLEX);
Complex *b = (Complex*) B.data;
for (int i=0; i<B.size; ++i) b[i] = g(b[i]);
lunum_pusharray1(L, &B);
}
}
static int luaC_lunum_fromfile(lua_State *L)
// -----------------------------------------------------------------------------
// Opens the binary file 'fname' for reading, and returns a 1d array from the
// data. The file size must be a multiple of the data type 'T'.
// -----------------------------------------------------------------------------
{
const char *fname = luaL_checkstring(L, 1);
const ArrayType T = luaL_optinteger(L, 2, ARRAY_TYPE_DOUBLE);
const int sizeof_T = array_sizeof(T);
FILE *input = fopen(fname, "rb");
if (input == NULL) {
return luaL_error(L, "no such file %s", fname);
}
fseek(input, 0L, SEEK_END); const size_t sz = ftell(input);
fseek(input, 0L, SEEK_SET);
if (sz % sizeof_T != 0) {
return luaL_error(L, "file size must be a multiple of the data type size");
}
const size_t N = sz / sizeof_T;
Array A = array_new_zeros(N, T);
if (fread(A.data, N, sizeof_T, input) != sizeof_T) {
fclose(input);
return luaL_error(L, "Error while reading file %s", fname);
}
fclose(input);
lunum_pusharray1(L, &A);
return 1;
}
static int _array_array_binary_op(lua_State *L, ArrayBinaryOperation op)
{
const Array *A = lunum_checkarray1(L, 1);
const Array *B = lunum_checkarray1(L, 2);
/* check size and dimensions */
if (A->ndims != B->ndims) {
return luaL_error(L, "arrays have different dimensions");
}
for (int d=0; d<A->ndims; ++d) {
if (A->shape[d] != B->shape[d]) {
return luaL_error(L, "arrays shapes do not agree");
}
}
const ArrayType T = (A->dtype >= B->dtype) ? A->dtype : B->dtype;
Array C = array_new_zeros(A->size, T);
array_resize_t(&C, A->shape, A->ndims);
lunum_pusharray1(L, &C);
array_array_binary_op(L, A, B, &C, op);
return 1;
}
static int luaC_lunum_slice(lua_State *L)
{
// The first part of this function extracts a slice of the array 'A' according
// to the convention start:stop:skip. The result is a contiguous array 'B'
// having the same number of dimensions as 'A'.
// ---------------------------------------------------------------------------
size_t Nd0_t, Nd1_t, Nd2_t, Nd3_t;
const Array *A = lunum_checkarray1(L, 1); // the array to resize
size_t *start = (size_t*) lunum_checkarray2(L, 2, ARRAY_TYPE_SIZE_T, &Nd0_t);
size_t *stop = (size_t*) lunum_checkarray2(L, 3, ARRAY_TYPE_SIZE_T, &Nd1_t);
size_t *skip = (size_t*) lunum_checkarray2(L, 4, ARRAY_TYPE_SIZE_T, &Nd2_t);
size_t *squeeze = (size_t*) lunum_checkarray2(L, 5, ARRAY_TYPE_SIZE_T, &Nd3_t);
int Nd0 = (int)Nd0_t, Nd1 = (int)Nd1_t, Nd2 = (int)Nd2_t, Nd3 = (int)Nd3_t;
if (Nd0 != A->ndims || Nd1 != A->ndims || Nd2 != A->ndims || Nd3 != A->ndims) {
return luaL_error(L, "slice has wrong number of dimensions for array");
}
for (int d=0; d<A->ndims; ++d) {
if (start[d] < 0 || stop[d] > A->shape[d]) {
return luaL_error(L, "slice not within array extent");
}
}
Array B = array_new_from_slice(A, start, stop, skip, Nd0);
// The rest of this function deals with squeezing out the size-1 dimensions of
// 'B' which are marked by the 'squeeze' array.
// ---------------------------------------------------------------------------
size_t Nd_new = 0;
for (int d=0; d<Nd0; ++d) Nd_new += !squeeze[d];
// In case we're left with a 0-dimensional (scalar) slice
if (Nd_new == 0) {
_push_value(L, B.dtype, B.data);
return 1;
}
// In case there are any dims to squeeze out
else if (Nd_new != Nd0) {
size_t *shape_new = (size_t*) malloc(Nd_new * sizeof(size_t));
for (int d=0,e=0; d<Nd0; ++d) {
if (B.shape[d] > 1 || !squeeze[d]) {
shape_new[e] = B.shape[d];
++e;
}
}
array_resize_t(&B, shape_new, Nd_new);
free(shape_new);
}
lunum_pusharray1(L, &B);
return 1;
}
int luaC_lunum_range(lua_State *L)
{
const int N = luaL_checkinteger(L, 1);
struct Array A = array_new_zeros(N, ARRAY_TYPE_INT);
lunum_pusharray1(L, &A);
for (int i=0; i<N; ++i) {
((int*)A.data)[i] = i;
}
return 1;
}
int test_astable(lua_State *L)
{
size_t N = 100;
Array A = array_new_zeros(N, ARRAY_TYPE_DOUBLE);
lunum_pusharray1(L, &A);
lunum_astable(L, -1);
lua_replace(L, -2);
return 1;
}
static int luaC_lunum_range(lua_State *L)
{
const lua_Integer N = luaL_checkinteger(L, 1);
if (N <= 0) { return luaL_error(L, "Invalid size %d", N); }
if (N <= INT_MAX) {
Array A = array_new_zeros(N, ARRAY_TYPE_INT);
lunum_pusharray1(L, &A);
for (size_t i=0; i<N; ++i) {
((int*)A.data)[i] = i;
}
} else {
Array A = array_new_zeros(N, ARRAY_TYPE_LONG);
lunum_pusharray1(L, &A);
for (size_t i=0; i<N; ++i) {
((long*)A.data)[i] = i;
}
}
return 1;
}
static int luaC_lunum_transpose(lua_State *L)
{
const Array *A = lunum_checkarray1(L, 1); // the array to transpose
/* copy transposed shape and size to B */
Array B = array_new_zeros(A->size, A->dtype);
array_resize_t(&B, A->shape, A->ndims);
array_transpose(A, &B);
lunum_pusharray1(L, &B);
return 1;
}
static int _array_unary_op(lua_State *L, ArrayUnaryOperation op)
{
Array *A = lunum_checkarray1(L, 1);
const size_t N = A->size;
ArrayType T = A->dtype;
Array B = array_new_zeros(N, T);
array_resize_t(&B, A->shape, A->ndims);
lunum_pusharray1(L, &B);
array_unary_op(L, A, &B, op);
return 1;
}
int luaC_array_astype(lua_State *L)
{
struct Array *A = lunum_checkarray1(L, 1);
enum ArrayType T;
if (lua_type(L, 2) == LUA_TSTRING) {
T = array_typeflag(lua_tostring(L, 2)[0]);
}
else {
T = (enum ArrayType) luaL_checkinteger(L, 2);
}
struct Array B = array_new_copy(A, T);
lunum_pusharray1(L, &B);
return 1;
}
static int luaC_lunum_linear(lua_State *L)
{
const lua_Number e1 = luaL_checknumber(L, 1);
const lua_Number e2 = luaL_checknumber(L, 2);
const lua_Integer N = luaL_checkinteger(L, 3);
if (N <= 1) { return luaL_error(L, "Invalid size %d", N); }
const ArrayType T = (ArrayType) luaL_optinteger(L, 4, ARRAY_TYPE_DOUBLE);
Array A = array_new_zeros(N, T);
void *a = A.data;
ARRAY_ASSIGN_OP(T, EXPR_ASSIGN_LINEAR);
lunum_pusharray1(L, &A);
return 1;
}
int luaC_array__index(lua_State *L)
{
struct Array *A = lunum_checkarray1(L, 1);
// Figure out what is the format of the input index. If it's a number or a
// table of numbers, then pass it along to _get_index. If it's a table of
// tables or numbers, then assume it's a slice. If it's an array of bools,
// then use it as a mask.
// ---------------------------------------------------------------------------
if (lunum_hasmetatable(L, 2, "array")) {
struct Array *M = lunum_checkarray1(L, 2);
if (M->dtype != ARRAY_TYPE_BOOL) {
luaL_error(L, "index array must be of type bool");
}
struct Array B = array_new_from_mask(A, M);
lunum_pusharray1(L, &B);
return 1;
}
else if (lua_type(L, 2) == LUA_TTABLE || lua_type(L, 2) == LUA_TSTRING) {
lua_getglobal(L, "lunum");
lua_getfield(L, -1, "__build_slice");
lua_remove(L, -2);
lua_pushvalue(L, 1);
lua_pushvalue(L, 2);
lua_call(L, 2, 1);
return 1;
}
const int m = _get_index(L, A);
_push_value(L, A->dtype, (char*)A->data + array_sizeof(A->dtype)*m);
return 1;
}
static int luaC_array__index(lua_State *L)
{
Array *A = lunum_checkarray1(L, 1);
// Figure out what is the format of the input index. If it's a number or a
// table of numbers, then pass it along to _get_index. If it's an array of bools,
// then use it as a mask.
// ---------------------------------------------------------------------------
if (lunum_hasmetatable(L, 2, "array")) {
Array *M = lunum_checkarray1(L, 2);
if (M->dtype != ARRAY_TYPE_BOOL) {
return luaL_error(L, "index array must be of type bool");
}
Array B = array_new_from_mask(A, M);
lunum_pusharray1(L, &B);
return 1;
}
/* try to index into array */
int success;
const size_t m = _get_index(L, A, &success);
if (success) {
_push_value(L, A->dtype, (char*)A->data + array_sizeof(A->dtype)*m);
return 1;
}
/* check metatable */
lua_getmetatable(L, 1);
lua_pushvalue(L, 2);
if (lua_gettable(L, -2) != LUA_TNIL) {
return 1;
}
return 0;
}
int _array_binary_op2(lua_State *L, enum ArrayOperation op)
{
struct Array *A = lunum_checkarray1(L, 1);
struct Array *B = lunum_checkarray1(L, 2);
if (A->ndims != B->ndims) {
luaL_error(L, "arrays have different dimensions");
}
for (int d=0; d<A->ndims; ++d) {
if (A->shape[d] != B->shape[d]) {
luaL_error(L, "arrays shapes do not agree");
}
}
const int N = A->size;
enum ArrayType T = (A->dtype >= B->dtype) ? A->dtype : B->dtype;
struct Array A_ = (A->dtype == T) ? *A : array_new_copy(A, T);
struct Array B_ = (B->dtype == T) ? *B : array_new_copy(B, T);
struct Array C = array_new_zeros(N, T);
array_resize(&C, A->shape, A->ndims);
lunum_pusharray1(L, &C);
array_binary_op(&A_, &B_, &C, op);
luaL_getmetatable(L, "array");
lua_setmetatable(L, -2);
if (A->dtype != T) array_del(&A_);
if (B->dtype != T) array_del(&B_);
return 1;
}
int lunum_upcast(lua_State *L, int pos, enum ArrayType T, int N)
// -----------------------------------------------------------------------------
// If the object at position 'pos' is already an array of dtype 'T', then push
// nothing and return 0. If the dtype is not 'T', then return 1 and push a copy
// of that array with dtype 'T' onto the stack. If it is a table, then push an
// array of dtype 'T' having the length of the table. If it is a number or
// complex, then push an array of dtype float or complex respectively having
// length 'N'.
// -----------------------------------------------------------------------------
{
if (array_typename(T) == NULL) {
luaL_error(L, "invalid array type");
}
// Deal with lunum.array
// ---------------------------------------------------------------------------
if (lunum_hasmetatable(L, pos, "array")) {
struct Array *A = lunum_checkarray1(L, pos);
if (A->dtype == T) {
return 0;
}
else {
struct Array A_ = array_new_copy(A, T);
lunum_pusharray1(L, &A_);
return 1;
}
}
// Deal with Lua table
// ---------------------------------------------------------------------------
else if (lua_istable(L, pos)) {
struct Array A = array_new_zeros(lua_rawlen(L, pos), T);
for (int i=0; i<A.size; ++i) {
lua_pushnumber(L, i+1);
lua_gettable(L, pos);
void *val = lunum_tovalue(L, T);
memcpy((char*)A.data + array_sizeof(T)*i, val, array_sizeof(T));
free(val);
lua_pop(L, 1);
}
lunum_pusharray1(L, &A);
return 1;
}
// Deal with Lua bool
// ---------------------------------------------------------------------------
else if (lua_isboolean(L, pos)) {
const Bool x = lua_toboolean(L, pos);
struct Array A = array_new_zeros(N, ARRAY_TYPE_BOOL);
array_assign_from_scalar(&A, &x);
lunum_pusharray1(L, &A);
return 1;
}
// Deal with Lua numbers
// ---------------------------------------------------------------------------
else if (lua_isnumber(L, pos)) {
const double x = lua_tonumber(L, pos);
struct Array A = array_new_zeros(N, ARRAY_TYPE_DOUBLE);
array_assign_from_scalar(&A, &x);
struct Array B = array_new_copy(&A, T);
array_del(&A);
lunum_pusharray1(L, &B);
return 1;
}
// Deal with lunum.complex
// ---------------------------------------------------------------------------
else if (lunum_hasmetatable(L, pos, "complex")) {
const Complex z = *((Complex*) lua_touserdata(L, pos));
struct Array A = array_new_zeros(N, ARRAY_TYPE_COMPLEX);
array_assign_from_scalar(&A, &z);
lunum_pusharray1(L, &A);
return 1;
}
// Throw an error
// ---------------------------------------------------------------------------
else {
luaL_error(L, "cannot cast to array from object of dtype %s\n",
lua_typename(L, lua_type(L, pos)));
return 0;
}
}
void lunum_pusharray2(lua_State *L, void *data, enum ArrayType T, int N)
{
struct Array A = array_new_zeros(N, T);
memcpy(A.data, data, N*array_sizeof(T));
lunum_pusharray1(L, &A);
}
static int _array_number_binary_op(lua_State *L, ArrayBinaryOperation op, Bool array_first)
{
const Array *A = array_first ? lunum_checkarray1(L, 1) : lunum_checkarray1(L, 2);
ArrayType T = A->dtype;
int num_pos = array_first ? 2 : 1;
union {
Bool b;
char c;
short s;
int i;
long l;
size_t t;
float f;
double d;
Complex z;
lua_Integer li;
lua_Number ln;
} num;
/* to force integer conversion if possible */
int isnum;
if (lua_isboolean(L, num_pos)) {
num.i = lua_toboolean(L, num_pos);
/* number can't have a higher type, upgrade to type T */
switch (T) {
case ARRAY_TYPE_BOOL : num.b = (Bool)num.i; break;
case ARRAY_TYPE_CHAR : num.c = (char)num.i; break;
case ARRAY_TYPE_SHORT : num.s = (short)num.i; break;
case ARRAY_TYPE_INT : num.i = (int)num.i; break;
case ARRAY_TYPE_LONG : num.l = (long)num.i; break;
case ARRAY_TYPE_SIZE_T : num.t = (size_t)num.i; break;
case ARRAY_TYPE_FLOAT : num.f = (float)num.i; break;
case ARRAY_TYPE_DOUBLE : num.d = (double)num.i; break;
case ARRAY_TYPE_COMPLEX : num.z = (Complex)num.i; break;
}
}
else if (num.li = lua_tointegerx(L, num_pos, &isnum), isnum) {
/* already assigned above */
if (T >= ARRAY_TYPE_LONG) {
/* A has higher type */
} else {
/* number has higher type */
T = ARRAY_TYPE_LONG;
}
/* upgrade to type T */
switch (T) {
case ARRAY_TYPE_BOOL : num.b = (Bool)num.li; break;
case ARRAY_TYPE_CHAR : num.c = (char)num.li; break;
case ARRAY_TYPE_SHORT : num.s = (short)num.li; break;
case ARRAY_TYPE_INT : num.i = (int)num.li; break;
case ARRAY_TYPE_LONG : num.l = (long)num.li; break;
case ARRAY_TYPE_SIZE_T : num.t = (size_t)num.li; break;
case ARRAY_TYPE_FLOAT : num.f = (float)num.li; break;
case ARRAY_TYPE_DOUBLE : num.d = (double)num.li; break;
case ARRAY_TYPE_COMPLEX : num.z = (Complex)num.li; break;
}
}
else if (num.ln = lua_tonumberx(L, num_pos, &isnum), isnum) {
/* already assigned above */
if (T >= ARRAY_TYPE_DOUBLE) {
/* A has higher type */
} else {
/* number has higher type */
T = ARRAY_TYPE_DOUBLE;
}
/* upgrade to type T */
switch (T) {
case ARRAY_TYPE_BOOL : num.b = (Bool)num.ln; break;
case ARRAY_TYPE_CHAR : num.c = (char)num.ln; break;
case ARRAY_TYPE_SHORT : num.s = (short)num.ln; break;
case ARRAY_TYPE_INT : num.i = (int)num.ln; break;
case ARRAY_TYPE_LONG : num.l = (long)num.ln; break;
case ARRAY_TYPE_SIZE_T : num.t = (size_t)num.ln; break;
case ARRAY_TYPE_FLOAT : num.f = (float)num.ln; break;
case ARRAY_TYPE_DOUBLE : num.d = (double)num.ln; break;
case ARRAY_TYPE_COMPLEX : num.z = (Complex)num.ln; break;
}
}
else if (lunum_hasmetatable(L, num_pos, "complex")) {
/* number complex */
num.z = *((Complex*) lua_touserdata(L, num_pos));
T = ARRAY_TYPE_COMPLEX;
} else {
return luaL_error(L, "Invalid argument in Array binary op");
}
Array C = array_new_zeros(A->size, T);
array_resize_t(&C, A->shape, A->ndims);
lunum_pusharray1(L, &C);
array_number_binary_op(L, A, (void *)&num, &C, op, array_first);
return 1;
}