/* methods */
static PyObject*
libtfr_tfr_spec(PyObject *self, PyObject *args)
{
/* arguments */
PyObject *o = NULL;
PyArrayObject *signal = NULL;
PyArrayObject *signal_cast = NULL;
PyObject *o2 = NULL;
PyArrayObject *fgrid = NULL;
PyArrayObject *fgrid_cast = NULL;
int N;
int step;
int Np;
int K = 6;
double tm = 6.0;
double flock = 0.01;
int tlock = 5;
/* output data */
npy_intp out_shape[2];
PyArrayObject *outdata = NULL;
double *spec;
/* internal stuff */
mfft *mtmh;
double *samples;
double *fgridp = NULL;
/* parse arguments */
if (!PyArg_ParseTuple(args, "Oiii|iddiO", &o, &N, &step, &Np, &K, &tm, &flock, &tlock, &o2))
return NULL;
signal = (PyArrayObject*) PyArray_FromAny(o, NULL, 1, 1, NPY_CONTIGUOUS, NULL);
if (signal==NULL) {
PyErr_SetString(PyExc_TypeError, "Input signal must be an ndarray");
return NULL;
}
int nfreq = N/2+1;
if (o2!=NULL && o2!=Py_None) {
fgrid = (PyArrayObject*) PyArray_FromAny(o2, NULL, 1, 1, NPY_CONTIGUOUS, NULL);
if (fgrid!=NULL) {
fgridp = coerce_ndarray_double(fgrid, &fgrid_cast);
if (fgridp==NULL) {
PyErr_SetString(PyExc_TypeError, "Unable to cast frequency grid to supported data type");
goto fail;
}
nfreq = PyArray_SIZE(fgrid);
}
}
/* coerce input data to proper type */
samples = coerce_ndarray_double(signal, &signal_cast);
if (samples==NULL) {
PyErr_SetString(PyExc_TypeError, "Unable to cast signal to supported data type");
goto fail;
}
/* initialize transform object - window size may be adjusted */
mtmh = mtm_init_herm(N, Np, K, tm);
/* allocate output array */
out_shape[0] = nfreq;
out_shape[1] = SPEC_NFRAMES(mtmh,PyArray_SIZE(signal),step);
outdata = (PyArrayObject*) PyArray_ZEROS(2,out_shape,NPY_DOUBLE,1); // last arg give fortran-order
spec = (double*) PyArray_DATA(outdata);
/* do the transform */
tfr_spec(mtmh, spec, samples, PyArray_SIZE(signal),
-1, step, flock, tlock, nfreq, fgridp);
mtm_destroy(mtmh);
Py_DECREF(signal);
Py_XDECREF(signal_cast);
Py_XDECREF(fgrid);
Py_XDECREF(fgrid_cast);
return PyArray_Return(outdata);
fail:
Py_XDECREF(fgrid_cast);
Py_XDECREF(fgrid);
Py_XDECREF(signal_cast);
Py_XDECREF(signal);
return NULL;
}
static PyObject *
potential_energy_and_forces(potential_t *self, PyObject *args, PyObject *kwargs)
{
static char *kwlist[] = {
"particles",
"neighbors",
"epot_per_at",
"epot_per_bond",
"f_per_bond",
"wpot_per_at",
"wpot_per_bond",
NULL
};
npy_intp dims[3];
npy_intp strides[3];
particles_t *a;
neighbors_t *n;
PyObject *return_epot_per_at = NULL;
PyObject *return_epot_per_bond = NULL;
PyObject *return_f_per_bond = NULL;
PyObject *return_wpot_per_at = NULL;
PyObject *return_wpot_per_bond = NULL;
int ierror = ERROR_NONE;
double epot;
PyObject *f;
PyObject *wpot;
PyObject *epot_per_at = NULL;
PyObject *epot_per_bond = NULL;
PyObject *f_per_bond = NULL;
PyObject *wpot_per_at = NULL;
PyObject *wpot_per_bond = NULL;
double *epot_per_at_ptr = NULL;
double *epot_per_bond_ptr = NULL;
double *f_per_bond_ptr = NULL;
double *wpot_per_at_ptr = NULL;
double *wpot_per_bond_ptr = NULL;
PyObject *r;
int i;
/* --- */
#ifdef DEBUG
printf("[potential_energy_and_forces] self = %p\n", self);
#endif
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!O!|O!O!O!O!O!", kwlist,
&particles_type, &a, &neighbors_type, &n,
&PyBool_Type, &return_epot_per_at,
&PyBool_Type, &return_epot_per_bond,
&PyBool_Type, &return_f_per_bond,
&PyBool_Type, &return_wpot_per_at,
&PyBool_Type, &return_wpot_per_bond))
return NULL;
epot = 0.0;
dims[0] = data_get_len(a->f90data);
dims[1] = 3;
strides[0] = dims[1]*NPY_SIZEOF_DOUBLE;
strides[1] = NPY_SIZEOF_DOUBLE;
f = (PyObject*) PyArray_New(&PyArray_Type, 2, dims, NPY_DOUBLE, strides,
NULL, 0, NPY_FARRAY, NULL);
memset(PyArray_DATA(f), 0, dims[0]*dims[1]*NPY_SIZEOF_DOUBLE);
dims[0] = 3;
dims[1] = 3;
wpot = PyArray_ZEROS(2, dims, NPY_DOUBLE, 1);
if (return_epot_per_at) {
if (return_epot_per_at == Py_True) {
dims[0] = data_get_len(a->f90data);
epot_per_at = PyArray_ZEROS(1, dims, NPY_DOUBLE, 1);
epot_per_at_ptr = PyArray_DATA(epot_per_at);
} else {
epot_per_at = Py_None;
Py_INCREF(Py_None);
}
}
if (return_epot_per_bond) {
if (return_epot_per_bond == Py_True) {
dims[0] = get_neighbors_size(n, a);
epot_per_bond = PyArray_ZEROS(1, dims, NPY_DOUBLE, 1);
epot_per_bond_ptr = PyArray_DATA(epot_per_bond);
} else {
epot_per_bond = Py_None;
Py_INCREF(Py_None);
}
}
if (return_f_per_bond) {
if (return_f_per_bond == Py_True) {
dims[0] = get_neighbors_size(n, a);
dims[1] = 3;
strides[0] = dims[1]*NPY_SIZEOF_DOUBLE;
strides[1] = NPY_SIZEOF_DOUBLE;
f_per_bond = (PyObject*) PyArray_New(&PyArray_Type, 2, dims,
NPY_DOUBLE, strides, NULL, 0,
NPY_FARRAY, NULL);
f_per_bond_ptr = PyArray_DATA(f_per_bond);
memset(f_per_bond_ptr, 0, dims[0]*dims[1]*NPY_SIZEOF_DOUBLE);
} else {
f_per_bond = Py_None;
Py_INCREF(Py_None);
}
}
if (return_wpot_per_at) {
if (return_wpot_per_at == Py_True) {
dims[0] = data_get_len(a->f90data);
dims[1] = 3;
dims[2] = 3;
strides[0] = dims[1]*dims[2]*NPY_SIZEOF_DOUBLE;
strides[1] = dims[2]*NPY_SIZEOF_DOUBLE;
strides[2] = NPY_SIZEOF_DOUBLE;
wpot_per_at = (PyObject*) PyArray_New(&PyArray_Type, 3, dims,
NPY_DOUBLE, strides, NULL, 0,
NPY_FARRAY, NULL);
wpot_per_at_ptr = PyArray_DATA(wpot_per_at);
memset(wpot_per_at_ptr, 0, dims[0]*dims[1]*dims[2]*NPY_SIZEOF_DOUBLE);
} else {
wpot_per_at = Py_None;
Py_INCREF(Py_None);
}
}
if (return_wpot_per_bond) {
if (return_wpot_per_bond == Py_True) {
dims[0] = get_neighbors_size(n, a);
dims[1] = 3;
dims[2] = 3;
strides[0] = dims[1]*dims[2]*NPY_SIZEOF_DOUBLE;
strides[1] = dims[2]*NPY_SIZEOF_DOUBLE;
strides[2] = NPY_SIZEOF_DOUBLE;
wpot_per_bond = (PyObject*) PyArray_New(&PyArray_Type, 3, dims,
NPY_DOUBLE, strides, NULL,
0, NPY_FARRAY, NULL);
wpot_per_bond_ptr = PyArray_DATA(wpot_per_bond);
memset(wpot_per_bond_ptr, 0, dims[0]*dims[1]*dims[2]*NPY_SIZEOF_DOUBLE);
} else {
wpot_per_bond = Py_None;
Py_INCREF(Py_None);
}
}
#ifdef DEBUG
printf("[potential_energy_and_forces] self->f90class->name = %s\n",
self->f90class->name);
printf("[potential_energy_and_forces] self->f90obj = %p\n",
self->f90obj);
printf("[potential_energy_and_forces] a->f90obj = %p\n",
a->f90obj);
printf("[potential_energy_and_forces] n->f90obj = %p\n",
n->f90obj);
printf("[potential_energy_and_forces] self->f90class->energy_and_forces = %p\n",
self->f90class->energy_and_forces);
#endif
self->f90class->energy_and_forces(self->f90obj, a->f90obj, n->f90obj,
&epot, PyArray_DATA(f), PyArray_DATA(wpot),
epot_per_at_ptr, epot_per_bond_ptr,
f_per_bond_ptr, wpot_per_at_ptr,
wpot_per_bond_ptr,
&ierror);
/*
* Now we need to reorder the per-bond properties such that some Python
* script can actually make sense out of the data.
*/
if (epot_per_bond_ptr) {
dims[0] = get_number_of_all_neighbors(n, a);
PyObject *tmp = PyArray_ZEROS(1, dims, NPY_DOUBLE, 1);
f_pack_per_bond_scalar(n->f90obj, epot_per_bond_ptr, PyArray_DATA(tmp));
Py_DECREF(epot_per_bond);
epot_per_bond = tmp;
}
if (wpot_per_bond_ptr) {
dims[0] = get_number_of_all_neighbors(n, a);
dims[1] = 3;
dims[2] = 3;
PyObject *tmp = PyArray_ZEROS(3, dims, NPY_DOUBLE, 0);
f_pack_per_bond_3x3(n->f90obj, wpot_per_bond_ptr, PyArray_DATA(tmp));
Py_DECREF(wpot_per_bond);
wpot_per_bond = tmp;
}
#ifdef DEBUG
printf("[potential_energy_and_forces] epot = %f\n", epot);
#endif
if (error_to_py(ierror))
return NULL;
/* --- Compose return tuple --- */
i = 3;
if (epot_per_at) i++;
if (epot_per_bond) i++;
if (f_per_bond) i++;
if (wpot_per_at) i++;
if (wpot_per_bond) i++;
r = PyTuple_New(i);
if (!r)
return NULL;
PyTuple_SET_ITEM(r, 0, PyFloat_FromDouble(epot));
PyTuple_SET_ITEM(r, 1, f);
PyTuple_SET_ITEM(r, 2, wpot);
i = 2;
if (epot_per_at) {
i++;
PyTuple_SET_ITEM(r, i, epot_per_at);
}
if (epot_per_bond) {
i++;
PyTuple_SET_ITEM(r, i, epot_per_bond);
}
if (f_per_bond) {
i++;
PyTuple_SET_ITEM(r, i, f_per_bond);
}
if (wpot_per_at) {
i++;
PyTuple_SET_ITEM(r, i, wpot_per_at);
}
if (wpot_per_bond) {
i++;
PyTuple_SET_ITEM(r, i, wpot_per_bond);
}
#ifdef DEBUG
printf("{potential_energy_and_forces}\n");
#endif
return r;
}
static PyObject*
libtfr_stft(PyObject *self, PyObject *args)
{
/* arguments */
PyObject *o = NULL;
PyArrayObject *signal = NULL;
PyArrayObject *signal_cast = NULL;
PyObject *o2 = NULL;
PyArrayObject *window = NULL;
int step;
int N = 0;
int Npoints, Ntapers;
/* output data */
npy_intp out_shape[3];
PyArrayObject *outdata = NULL;
int do_complex = 0;
double *spec;
complex double *zspec_tmp;
npy_cdouble *zspec;
/* internal stuff */
mfft *mtmh;
double *samples = NULL;
double *windowp;
/* parse arguments */
if (!PyArg_ParseTuple(args, "OOi|ii", &o, &o2, &step, &N, &do_complex))
return NULL;
signal = (PyArrayObject*) PyArray_FromAny(o, NULL, 1, 1, NPY_CONTIGUOUS, NULL);
if (signal==NULL) {
PyErr_SetString(PyExc_TypeError, "Input signal must be an ndarray");
return NULL;
}
window = (PyArrayObject*) PyArray_FromAny(o2, NULL, 1, 2, NPY_CONTIGUOUS, NULL);
if (window==NULL) {
PyErr_SetString(PyExc_TypeError, "Window must be a 1D or 2D ndarray");
goto fail;
}
/* determine dimensions of window */
if (PyArray_NDIM(window)==1) {
Ntapers = 1;
Npoints = PyArray_DIM(window, 0);
}
else {
Ntapers = PyArray_DIM(window, 0);
Npoints = PyArray_DIM(window, 1);
}
/* coerce data to proper type */
samples = coerce_ndarray_double(signal, &signal_cast);
if (samples==NULL) {
PyErr_SetString(PyExc_TypeError, "Unable to cast signal to supported data type");
goto fail;
}
/* need to copy the window function b/c mtm_destroy() will demalloc it */
if (PyArray_TYPE(window)!=NPY_DOUBLE) {
PyErr_SetString(PyExc_TypeError, "Window function must be double precision float");
goto fail;
}
windowp = malloc(Npoints * Ntapers * sizeof(double));
memcpy(windowp, PyArray_DATA(window), Npoints * Ntapers * sizeof(double));
/* allocate outputs and do the transform */
if (N < 1)
N = Npoints;
mtmh = mtm_init(N, Npoints, Ntapers, windowp, NULL);
out_shape[0] = N/2+1;
if (do_complex) {
out_shape[1] = Ntapers;
out_shape[2] = SPEC_NFRAMES(mtmh, PyArray_SIZE(signal), step);
outdata = (PyArrayObject*) PyArray_ZEROS(3,out_shape,NPY_CDOUBLE,1); // fortran-order
zspec = (npy_cdouble*) PyArray_DATA(outdata);
//printf("output dimensions: %d, %d, %d\n", out_shape[0], out_shape[1], out_shape[2]);
zspec_tmp = (complex double*)malloc(PyArray_SIZE(outdata) * sizeof(complex double));
mtm_zspec(mtmh, zspec_tmp, samples, PyArray_SIZE(signal), step);
cmplx_c99tonpy(zspec_tmp, zspec, PyArray_SIZE(outdata));
free(zspec_tmp);
}
else {
out_shape[1] = SPEC_NFRAMES(mtmh, PyArray_SIZE(signal), step);
outdata = (PyArrayObject*) PyArray_ZEROS(2,out_shape,NPY_DOUBLE,1); // fortran-order
spec = (double*) PyArray_DATA(outdata);
mtm_spec(mtmh, spec, samples, PyArray_SIZE(signal), step, 0);
}
mtm_destroy(mtmh);
Py_DECREF(signal);
Py_DECREF(window);
Py_XDECREF(signal_cast);
return PyArray_Return(outdata);
fail:
Py_XDECREF(signal_cast);
Py_XDECREF(signal);
Py_XDECREF(window);
Py_XDECREF(outdata);
return NULL;
}
static PyObject*
transform (PyObject *dummy, PyObject *args)
{
PyObject *intsarg=NULL;
PyObject *orbitalsarg=NULL;
PyObject *ints=NULL;
PyObject *orbitals=NULL;
if (!PyArg_ParseTuple(args, "OO", &intsarg, &orbitalsarg))
return NULL;
ints = PyArray_FROM_OTF(intsarg, NPY_DOUBLE, NPY_IN_ARRAY);
if (ints == NULL)
return NULL;
int intsD = PyArray_NDIM(ints);
if (intsD != 4)
{
printf("First input varaible must be a four dimensional array\n");
return NULL;
}
int Nin = PyArray_DIM(ints, 0);
for(int i = 0; i < intsD; i++)
{
if(PyArray_DIM(ints, i) != Nin)
{
printf("First input varaible must have all dimensions of the same size\n");
return NULL;
}
}
orbitals = PyArray_FROM_OTF(orbitalsarg, NPY_DOUBLE, NPY_IN_ARRAY);
if (orbitals == NULL)
return NULL;
int orbitalsD = PyArray_NDIM(orbitals);
if (orbitalsD != 2)
{
printf("Second input varaible must be two dimensional array\n");
return NULL;
}
if (Nin != PyArray_DIM(orbitals, 0))
{
printf("Second input has wrong dimension. Does not fit to first varaible.\n");
return NULL;
}
int Nout = PyArray_DIM(orbitals, 1);
//printf("Loaded integrals for basis size %d x %d.\n", Nin, Nout);
npy_intp dims[4] = { Nout, Nout, Nout, Nout };
PyObject* result = PyArray_ZEROS(4, dims, NPY_DOUBLE, 0);
double* ints_data = PyArray_DATA(ints);
double* orbitals_data = PyArray_DATA(orbitals);
double* result_data = PyArray_DATA(result);
for (int i = 0; i < Nout; ++i)
{
for (int j = i; j < Nout; ++j)
{
printf("i = %d j = %d\n", i, j);
for (int k = i; k < Nout; ++k)
{
for (int l = k; l < Nout; ++l)
{
if (abs(result_data[i * Nout * Nout * Nout + j * Nout * Nout + k * Nout + l]) < 1e-5)
{
double sum = 0;
for (int mu = 0; mu < Nin; ++mu)
{
for (int nu = 0; nu < Nin; ++nu)
{
for (int lamb = 0; lamb < Nin; ++lamb)
{
for (int sigma = 0; sigma < Nin; ++sigma)
{
sum +=
orbitals_data[mu * Nout + i] *
orbitals_data[nu * Nout + j] *
ints_data[mu * Nin * Nin * Nin + nu * Nin * Nin + lamb * Nin + sigma] *
orbitals_data[lamb * Nout + k] *
orbitals_data[sigma * Nout + l];
}
}
}
}
result_data[i * Nout * Nout * Nout + j * Nout * Nout + k * Nout + l] = sum;
result_data[i * Nout * Nout * Nout + j * Nout * Nout + l * Nout + k] = sum;
result_data[j * Nout * Nout * Nout + i * Nout * Nout + k * Nout + l] = sum;
result_data[j * Nout * Nout * Nout + i * Nout * Nout + l * Nout + k] = sum;
result_data[k * Nout * Nout * Nout + l * Nout * Nout + i * Nout + j] = sum;
result_data[k * Nout * Nout * Nout + l * Nout * Nout + j * Nout + i] = sum;
result_data[l * Nout * Nout * Nout + k * Nout * Nout + i * Nout + j] = sum;
result_data[l * Nout * Nout * Nout + k * Nout * Nout + j * Nout + i] = sum;
}
}
}
}
}
Py_DECREF(ints);
Py_DECREF(orbitals);
return result;
}
