c/interaction/interaction.c

@@ -30,24 +30,14 @@ interaction_ufunc_force
 	, void * data)
 {
 	char * in = args[0]
-		, * raw_coefficients = args[1]
-		, * out = args[2];
+		, * out = args[1];
 	npy_intp n = dimensions[0];
 	npy_intp in_step = steps[0]
-		, raw_coefficients_steps = steps[1]
-		, out_step = steps[2];
+		, out_step = steps[1];
 
 	npy_intp i;
 
-	// Copy the coefficients to a safe array.
-	// This is because NumPy arrays are not 
-	// necessarily cast-safe.
-	float coefficients[19];
-	for(i = 0; i < 19; i++)
-	{
-		coefficients[i] = *(float *)raw_coefficients;
-		raw_coefficients += raw_coefficients_steps;
-	}
+	float * coefficients = (float *) data;
 
 	for(i = 0; i < n; i++)
 	{
@@ -73,27 +63,17 @@ interaction_ufunc_float2D
 		, * y_old = args[1]
 		, * p_x_old = args[2]
 		, * p_y_old = args[3]
-		, * raw_coefficients = args[4]
-		, * p_x_new = args[5]
-		, * p_y_new = args[6];
+		, * p_x_new = args[4]
+		, * p_y_new = args[5];
 
 	npy_intp x_old_steps = steps[0]
 		, y_old_steps = steps[1]
 		, p_x_old_steps = steps[2]
 		, p_y_old_steps = steps[3]
-		, raw_coefficients_steps = steps[4]
-		, p_x_new_steps = steps[5]
-		, p_y_new_steps = steps[6];
+		, p_x_new_steps = steps[4]
+		, p_y_new_steps = steps[5];
 
-	// Copy the coefficients to a safe array.
-	// This is because NumPy arrays are not 
-	// necessarily cast-safe.
-	float coefficients[19];
-	for(i = 0; i < 19; i++)
-	{
-		coefficients[i] = *(float *)raw_coefficients;
-		raw_coefficients += raw_coefficients_steps;
-	}
+	float *coefficients = (float *) data;
 
 	// Compute the new momenta: 
 
@@ -135,30 +115,159 @@ interaction_ufunc_float2D
 			// Update the momenta.
 			delta_p_x = delta_x_e * interaction_force_function(r, coefficients);
 			delta_p_y = delty_y_e * interaction_force_function(r, coefficients);
-			*(float *)(p_x_new + i*p_x_new_steps) -= delta_p_x;
-			*(float *)(p_y_new + i*p_y_new_steps) -= delta_p_y;
-			*(float *)(p_x_new + j*p_x_new_steps) += delta_p_x;
-			*(float *)(p_y_new + j*p_y_new_steps) += delta_p_y;
+			//printf("%d, %d: %f, %f\n", i, j, delta_p_x, delta_p_y);
+			*(float *)(p_x_new + i*p_x_new_steps) += delta_p_x;
+			*(float *)(p_y_new + i*p_y_new_steps) += delta_p_y;
+			*(float *)(p_x_new + j*p_x_new_steps) -= delta_p_x;
+			*(float *)(p_y_new + j*p_y_new_steps) -= delta_p_y;
 		}
 	}
 }
+static char interaction_types[] = 
+	{ NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT};
+static char force_types[] =
+	{ NPY_FLOAT, NPY_FLOAT};
+static PyUFuncGenericFunction force_funcs[1] = 
+	{ interaction_ufunc_force};
+static PyUFuncGenericFunction interaction_funcs[1] = 
+	{ interaction_ufunc_float2D};
+
+typedef struct
+{
+	PyObject_HEAD
+	float coefficients[19];
+	PyObject * ufunc;
+	void *data[1];
+} interaction_UFuncWrapper;
+
+static int
+interaction_UFuncWrapper_init
+	( interaction_UFuncWrapper * self
+	, PyObject * args
+	, PyObject * kwds)
+{
+	// 0: ufunc_force
+	// 1: interaction2D
+	char type;
+	PyObject * coefficients;
+	int i;
+	PyObject * this_coefficient;
+
+	char *kwords[] = { "type_", "coefficients", NULL};
+
+	if(!PyArg_ParseTupleAndKeywords(args, kwds, "bO", kwords, &type, &coefficients))
+	{
+		return -1;
+	}
+	if(!PySequence_Check(coefficients))
+	{
+		return -1;
+	}
+
+	if(PySequence_Size(coefficients) != 19)
+	{
+		PyErr_SetString(PyExc_ValueError, "coefficients must have length 19");
+		return -1;
+	}
+
+	// copy the coefficients.
+	for(i = 0; i < 19; i++)
+	{
+		this_coefficient = PySequence_GetItem(coefficients, i);
+		if(!this_coefficient)
+		{
+			return -1;
+		}
+		// XXX: PyFloat_AsDouble might call python code,
+		// so make sure that nothing bad can happen.
+		Py_INCREF(this_coefficient);
+		self->coefficients[i] = PyFloat_AsDouble(this_coefficient);
+		Py_DECREF(this_coefficient);
+		if(PyErr_Occurred())
+		{
+			return -1;
+		}
+	}
+	self->data[0] = (void *)self->coefficients;
+
+	switch(type)
+	{
+		case 0:
+		{
+			self->ufunc = PyUFunc_FromFuncAndData(
+				force_funcs // func
+				, self->data // data
+				, force_types //types
+				, 1 // ntypes
+				, 1 // nin
+				, 1 // nout
+				, PyUFunc_None // identity
+				, "force_function" // name
+				, "computes the scalar force between two particles with given coefficients" // doc
+				, 0); // unused
+			break;
+		}
+		case 1:
+		{
+			self->ufunc = PyUFunc_FromFuncAndData(
+				interaction_funcs
+				, self->data
+				, interaction_types
+				, 1
+				, 4
+				, 2
+				, PyUFunc_None
+				, "interaction2D"
+				, "Update the momenta according to the given coefficients and positions"
+				, 0);
+			break;
+		}
+		default:
+		{
+			PyErr_SetString(PyExc_ValueError, "unknown ufunc type, must be 0 or 1");
+			return -1;
+		}
+	}
+	Py_INCREF(self->ufunc);
+	return 0;
+}
+
+
+static PyObject *
+interaction_UFuncWrapper_call
+	(interaction_UFuncWrapper * self
+	 , PyObject * args
+	 , PyObject * kwargs)
+{
+	return PyObject_Call(self->ufunc, args, kwargs);
+}
+
+static PyMemberDef interaction_UFuncWrapper_members[] = 
+{
+	{"ufunc", T_OBJECT_EX, offsetof(interaction_UFuncWrapper, ufunc), 0, "ufunc"},
+	{NULL}
+};
+
+static PyTypeObject interaction_UFuncWrapperType = 
+{
+	PyVarObject_HEAD_INIT(NULL, 0)
+	.tp_name = "brown.interaction.UFuncWrapper",
+	.tp_doc = "A wrapper that wraps the ufuncs for interaction and force, storing the coefficients",
+	.tp_basicsize = sizeof(interaction_UFuncWrapper),
+	.tp_itemsize = 0,
+	.tp_flags = Py_TPFLAGS_DEFAULT,
+	.tp_new = PyType_GenericNew,
+	.tp_init = (initproc) interaction_UFuncWrapper_init,
+	.tp_call = interaction_UFuncWrapper_call,
+	.tp_members = interaction_UFuncWrapper_members
+};
+
+
 
 static PyMethodDef InteractionMethods[] = {
 	{NULL, NULL, 0, NULL}
 };
 
-PyUFuncGenericFunction interaction_funcs[] = 
-	{ &interaction_ufunc_float2D};
-PyUFuncGenericFunction force_funcs[] = 
-	{ &interaction_ufunc_force};
-
-static char interaction_types[] = 
-	{ NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT};
-static char force_types[] =
-	{ NPY_FLOAT, NPY_FLOAT, NPY_FLOAT};
-
-static void *interaction_data[] = {NULL};
-static void *force_data[] = {NULL};
 
 static struct PyModuleDef moduledef = {
 	PyModuleDef_HEAD_INIT
@@ -175,10 +284,12 @@ static struct PyModuleDef moduledef = {
 PyMODINIT_FUNC 
 PyInit_interaction(void)
 {
-	PyObject * module
-		, * ufunc_interaction
-		, * ufunc_force
-		, * dct;
+	PyObject * module;
+
+	if(PyType_Ready(&interaction_UFuncWrapperType) < 0)
+	{
+		return NULL;
+	}
 
 	module = PyModule_Create(&moduledef);
 	if(!module)
@@ -186,38 +297,10 @@ PyInit_interaction(void)
 		return NULL;
 	}
 	import_array();
-	import_umath();
+	import_ufunc();
 
-	ufunc_interaction = PyUFunc_FromFuncAndDataAndSignature(
-			interaction_funcs
-			, interaction_data
-			, interaction_types
-			, 1
-			, 5
-			, 2
-			, PyUFunc_None
-			, "interaction2D"
-			, "Update the momenta according to the given coefficients and positions"
-			, 0
-			, "(n),(n),(n),(n),(19)->(n),(n)");
-	ufunc_force = PyUFunc_FromFuncAndDataAndSignature(
-			force_funcs
-			, force_data
-			, force_types
-			, 1
-			, 2
-			, 1
-			, PyUFunc_None
-			, "force_function"
-			, "computes the scalar force between two particles with given coefficients"
-			, 0
-			, "(n),(19)->(n)");
-
-	dct = PyModule_GetDict(module);
-	PyDict_SetItemString(dct, "interaction2D", ufunc_interaction);
-	PyDict_SetItemString(dct, "force_function", ufunc_force);
-	Py_DECREF(ufunc_interaction);
-	Py_DECREF(ufunc_force);
+	Py_INCREF(&interaction_UFuncWrapperType);
+	PyModule_AddObject(module, "UFuncWrapper", (PyObject *) &interaction_UFuncWrapperType);
 
 	return module;
 }

c/interaction/interaction.h

@@ -2,8 +2,12 @@
 #define interaction_h
 
 #include <Python.h>
+#include "structmember.h"
+// XXX
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
 #include <numpy/ndarraytypes.h>
 #include <numpy/ufuncobject.h>
+#include <stddef.h>
 
 /*
  * This is a quite generic force function mapping a 
@@ -28,5 +32,11 @@ interaction_ufunc_float2D
 	, npy_intp * steps
 	, void * data);
 
+static void
+interaction_ufunc_force
+	( char ** args
+	, npy_intp * dimensions
+	, npy_intp * steps
+	, void * data);
 
 #endif

particles.py

@@ -0,0 +1,57 @@
+from brown.interaction import UFuncWrapper
+from brown.brown import BrownIterator
+import numpy as np
+from collections import deque
+from copy import copy
+import matplotlib.pyplot as plt
+import matplotlib.animation as ani
+
+c = np.array([5, 10, 20, 30, 0, 0, 0, 1, -20, 0, -2, -0.1, 2, 0, 0, 0, 0, 0, 0], dtype=np.float16)
+#force_function = UFuncWrapper(0, c)
+#interaction2D = UFuncWrapper(1, c)
+
+borders_x = [-100, 100]
+borders_y = [-100, 100]
+n_particles = 6
+frames = 1000
+
+x_coords = np.random.uniform(borders_x[0] / 2, borders_x[1] / 2, n_particles).astype(np.float16)
+y_coords = np.random.uniform(borders_y[0] / 2, borders_y[1] / 2, n_particles).astype(np.float16)
+
+
+x_momenta = np.zeros(n_particles, dtype=np.float16)
+y_momenta = np.zeros(n_particles, dtype=np.float16)
+
+
+
+fig = plt.figure(figsize=(7, 7))
+ax = fig.add_axes([0, 0, 1, 1], frameon=False)
+ax.set_xlim(*borders_x)
+ax.set_xticks([])
+ax.set_ylim(*borders_y)
+ax.set_yticks([])
+
+plot, = ax.plot(x_coords, y_coords, "b.")
+center_of_mass, = ax.plot(x_coords.mean(), y_coords.mean(), "r-")
+center_of_mass_history_x = deque([x_coords.mean()])
+center_of_mass_history_y = deque([y_coords.mean()])
+
+brown = BrownIterator(-1, c
+			, x_coords, y_coords
+			, y_momenta, y_momenta
+			, borders_x, borders_y
+			, border_dampening=1
+			, dt=0.001)
+
+u = iter(brown)
+
+def update(i):
+	data = next(u)
+	center_of_mass_history_x.append(x_coords.mean())
+	center_of_mass_history_y.append(y_coords.mean())
+
+	plot.set_data(*data)
+	center_of_mass.set_data(center_of_mass_history_x, center_of_mass_history_y)
+
+animation = ani.FuncAnimation(fig, update, range(frames), interval=1)
+plt.show()

py/brown/brown.py

@@ -0,0 +1,66 @@
+from brown.interaction import UFuncWrapper
+import numpy as np
+
+class BrownIterator(object):
+	def __init__(self
+			, m
+			, c
+			, x
+			, y
+			, px
+			, py
+			, borders_x
+			, borders_y
+			, speed_of_light=1e3 # The value that will replace NaN momenta
+			, border_dampening=1
+			, dt=0.1):
+		self._max_m = m
+		self._i = 0
+		self.c = c
+		self.x = x
+		self.y = y
+		self.px = px
+		self.py = py
+		self.borders_x = borders_x
+		self.borders_y = borders_y
+		self.dt = dt
+		self.speed_of_light = speed_of_light
+		self.border_dampening = border_dampening
+		self._interaction = UFuncWrapper(1, c)
+	def __iter__(self):
+		self._i = 0
+		return self
+	def __next__(self):
+		self._i += 1
+		if(self._i > self._max_m and self._max_m > 0):
+			raise StopIteration()
+		if(self._i == 1):
+			return self.x, self.y
+
+		self.px, self.py = self._interaction(self.x, self.y, self.px, self.py)
+		self.px[np.isnan(self.px)] = self.speed_of_light
+		self.py[np.isnan(self.py)] = self.speed_of_light
+		self._reflect_at_borders()
+
+		self.x += self.dt * self.px
+		self.y += self.dt * self.py
+
+		return self.x, self.y
+
+
+	def _reflect_at_borders(self):
+		if(len(self.borders_x) > 0):
+			self.px[self.x <= self.borders_x[0]] *= -self.border_dampening
+			self.x[self.x <= self.borders_x[0]] = self.borders_x[0]
+		if(len(self.borders_x) > 1):
+			self.px[self.x >= self.borders_x[1]] *= -self.border_dampening
+			self.x[self.x >= self.borders_x[1]] = self.borders_x[1]
+		if(len(self.borders_y) > 0):
+			self.py[self.y <= self.borders_y[0]] *= -self.border_dampening
+			self.y[self.y <= self.borders_y[0]] = self.borders_y[0]
+		if(len(self.borders_y) > 1):
+			self.py[self.y >= self.borders_y[1]] *= -self.border_dampening
+			self.y[self.y >= self.borders_y[1]] = self.borders_y[1]
+
+		
+

setup.py

@@ -1,30 +1,31 @@
-#from distutils.core import setup, Extension
-#
-#
-#interaction = Extension("brown.interaction",
-#		sources = ["c/interaction/interaction.c"])
-#
-#setup(name = "brown",
-#	version = "0.0.1",
-#	description = "Me playing around with single-atom classical gases",
-#	ext_modules = [interaction],
-#	packages = [
-#		"brown"
-#	],
-#	package_dir = {"brown": "py/brown"},
-#	#url="https://github.com/daknuett/python3-nf",
-#	author = "Daniel Knüttel",
-#	author_email = "daniel.knuettel@daknuett.eu")
+from distutils.core import setup, Extension
 
-def configuration(parent_package='', top_path=None):
-	from numpy.distutils.misc_util import Configuration
 
-	config = Configuration('brown'
-			, parent_package
-			, top_path)
-	config.add_extension('interaction', ['c/interaction/interaction.c'])
-	return config
+interaction = Extension("brown.interaction",
+		sources = ["c/interaction/interaction.c"])
 
-if __name__ == "__main__":
-	from numpy.distutils.core import setup
-	setup(configuration=configuration)
+setup(name = "brown",
+	version = "0.0.1",
+	description = "Me playing around with single-atom classical gases",
+	ext_modules = [interaction],
+	packages = [
+		"brown"
+	],
+	package_dir = {"brown": "py/brown"},
+	#url="https://github.com/daknuett/python3-nf",
+	author = "Daniel Knüttel",
+	author_email = "daniel.knuettel@daknuett.eu")
+
+#def configuration(parent_package='', top_path=None):
+#	from numpy.distutils.misc_util import Configuration
+#
+#	config = Configuration('brown'
+#			, parent_package
+#			, top_path="py/")
+#	config.add_extension('interaction', ['c/interaction/interaction.c'], extra_compile_args=["-g"])
+#	#config.add_subpackage("brown.brown", "py/brown")
+#	return config
+#
+#if __name__ == "__main__":
+#	from numpy.distutils.core import setup
+#	setup(configuration=configuration)

test.py

@@ -1,46 +1,50 @@
-from brown.interaction import interaction2D,force_function
+from brown.interaction import UFuncWrapper
 import numpy as np
 from collections import deque
 from copy import copy
 import matplotlib.pyplot as plt
 
-c = np.array([5, 10, 20, 0, 0, 0, 0, 1, -2, 0, -2, -0.1, 2, 0, 0, 0, 0, 0, 0], dtype=np.float16)
-print(c)
+c = np.array([5, 10, 20, 30, 0, 0, 0, 1, -20, 0, -2, -0.1, 2, 0, 0, 0, 0, 0, 0], dtype=np.float16)
+force_function = UFuncWrapper(0, c)
+interaction2D = UFuncWrapper(1, c)
 
-#x_coords = np.array([0, 1], dtype=np.float16)
-#y_coords = np.array([0, 0], dtype=np.float16)
-#
-#x_momenta = np.array([0, 0], dtype=np.float16)
-#y_momenta = np.array([0, 0], dtype=np.float16)
-#
-#time = np.arange(0, 5, 1, dtype=np.float16)
-#
-#time_evolution = deque()
-#
-#for t in time:
-#	x_momenta, y_momenta = interaction2D(x_coords, y_coords, x_momenta, y_momenta, c)
-#	x_coords += x_momenta
-#	y_coords += y_momenta
-#
-#	time_evolution.append(copy(x_coords))
-#
-#time_evolution = np.array(time_evolution)
-#
-#particle_one_evolution = time_evolution[:, 0]
-#particle_two_evolution = time_evolution[:, 1]
+x_coords = np.array([0, 1], dtype=np.float16)
+y_coords = np.array([0, 0], dtype=np.float16)
 
-#plt.subplot(2, 1, 1)
-r = np.arange(0, 0.5, 0.01, dtype=np.float16)
+x_momenta = np.array([0, 0], dtype=np.float16)
+y_momenta = np.array([0, 0], dtype=np.float16)
+
+time = np.arange(0, 50, 1, dtype=np.float16)
+
+time_evolution = deque()
+
+dt = 0.1
+
+for t in time:
+	x_momenta, y_momenta = interaction2D(x_coords, y_coords, x_momenta, y_momenta)
+	x_coords += dt * x_momenta
+	y_coords += dt * y_momenta
+
+	time_evolution.append(copy(x_coords))
+
+time_evolution = np.array(time_evolution)
+
+particle_one_evolution = time_evolution[:, 0]
+particle_two_evolution = time_evolution[:, 1]
+
+plt.subplot(2, 1, 1)
+
+r = np.arange(0, 50, 0.1, dtype=np.float16)
 plt.title("Force")
 plt.xlabel("particle distance")
 plt.ylabel("scalar force")
-plt.plot(r, force_function(r, c))
+plt.plot(r, force_function(r))
 
-#plt.subplot(2, 1, 2)
-#plt.title("Particle x-positions over time")
-#plt.xlabel("time")
-#plt.ylabel("particle x-position")
-#h0, = plt.plot(time, particle_one_evolution, label="particle 1")
-#h1, = plt.plot(time, particle_two_evolution, label="particle 2")
-#plt.legend(handles=[h0, h1])
+plt.subplot(2, 1, 2)
+plt.title("Particle x-positions over time")
+plt.xlabel("time")
+plt.ylabel("particle x-position")
+h0, = plt.plot(time, particle_one_evolution, label="particle 1")
+h1, = plt.plot(time, particle_two_evolution, label="particle 2")
+plt.legend(handles=[h0, h1])
 plt.show()