We Keep Coding

Gaussian fit for Python with two different function

I used the Gaussian fit with 3 gauss to adjust but datai but I utility data that sometimes my curve contains only two Gaussians in it not find the parameter remnants to use and but great an error is what there is a method that but allows to change with curve fit function use if two or three gaussians .
for my function main, i have this code :
FitGWPS = mainCurveFitGWPS(global_ws, period, All_Max_GWPS, DoupleDip)
and my code for fit is :
import numpy as np
from scipy.optimize import curve_fit
#Functions-----------------------------------------
#Gaussian function
def _1gaus(X,C,X_mean,sigma):
return C*np.exp(-(X-X_mean)**2/(2*sigma**2))
def _3gaus(x, amp1,cen1,sigma1, amp2,cen2,sigma2, amp3,cen3,sigma3):
return amp1*np.exp(-(x-cen1)**2/(2*sigma1**2)) +\
amp2*np.exp(-(x-cen2)**2/(2*sigma2**2)) + amp3*np.exp(-(x-
cen3)**2/(2*sigma3**2))
def ParamFit (Gws, P, Max, popt_Firstgauss):
#Calculating the Lorentzian PDF values given Gaussian parameters and random variableX
width=0
Amp = []
cen = []
wid = []
for j in range(len(Max-1)):
Amp.append(0.8 * (Gws[Max[j]])) # Amplitude
cen.append(P[Max[j]]) # Frequency
if j == 0 : wid.append(0.3 + width * 2.) # Width
else : wid.append(0.3 + popt_Firstgauss[2] * 2.)
return Amp,wid,cen
def mainCurveFitGWPS(global_ws_in, period_in, All_Max_GWPS, DoupleDip):
#Calculating the Gaussian PDF values given Gaussian parameters and random variable X
# For the first fit we calculate with function of the max values
mean = sum(period_in*(global_ws_in))/sum((global_ws_in ))
sigma = np.sqrt(sum((global_ws_in)*(period_in-mean)**2)/sum((global_ws_in)))
Cst = 1 / ( 2* np.pi * sigma)
width=0
Amp = 0.8 * (global_ws_in[All_Max_GWPS[0]]) # Amplitude
cen = period_in[All_Max_GWPS[0]] # Frequency
wid = 0.3 + width * 2. #Width
Amp = []
cen = []
wid = []
for j in range(len(All_Max_GWPS-1)):
Amp.append(0.8 * (global_ws_in[All_Max_GWPS[j]])) # Amplitude
cen.append(period_in[All_Max_GWPS[j]]) # Frequency
if j == 0 : wid.append(0.3 + width * 2.)
else : wid.append(0.3 + popt_gauss[2] * 2.)
#do the fit!
popt_gauss, pcov_gauss = curve_fit(_1gaus, period_in, global_ws_in, p0 = [Cst,
mean, sigma])
FitGauss = _1gaus(period_in, *popt_gauss)
#I use the center, amplitude, and sigma values which I used to create the fake
#data
popt_3gauss, pcov_3gauss = curve_fit(_3gaus, period_in, global_ws_in, p0=[Amp[0],
cen[0], wid[0],Amp[1], cen[1], wid[1],Amp[2], cen[2], wid[2]], maxfev =5000)
Fit3Gauss = _3gaus(period_in, *popt_3gauss)
return Fit3Gauss
for example picture :
and

Normalise all vectors on the z axis

I have written code that calculates the angle between two vectors. However the way in which is does this is to start with two vectors, rotate each according to some euler angles calculated in a separate program, then calculate the angle between the vectors.
Up until now I have been working with a use case that means both starting vectors are (0,0,1) that makes life super easy. I could just take one set of euler angles away from the other and then calculate the angle between 0,0,1 and the vector that had been rotated by the difference. It meant I could plot nice distribution plots and vector diagrams because everything was normalised to 0,0,1. (I have 1000s of these vectors for the record).
No I am trying to write in a function that would allow for a use case where the two starting vectors are not on 0,0,1. I figured the easiest way to do this would be to calculate direction of the vector relative to 0,0,1 and after calculating the position of the vector just rotate by the precalculated offsets. (this might be a stupid way to do it, if it is please tell me).
MY current code works for a case where a vector is 0,1,0 but then breaks down if i start entering random numbers.
import numpy as np
import math
def RotationMatrix(axis, rotang):
"""
This uses Euler-Rodrigues formula.
"""
#Input taken in degrees, here we change it to radians
theta = rotang * 0.0174532925
axis = np.asarray(axis)
#Ensure axis is a unit vector
axis = axis/math.sqrt(np.dot(axis, axis))
#calclating a, b, c and d according to euler-rodrigues forumla requirments
a = math.cos(theta/2)
b, c, d = axis*math.sin(theta/2)
a2, b2, c2, d2 = a*a, b*b, c*c, d*d
bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
#Return the rotation matrix
return np.array([[a2+b2-c2-d2, 2*(bc-ad), 2*(bd+ac)],
[2*(bc+ad), a2+c2-b2-d2, 2*(cd-ab)],
[2*(bd-ac), 2*(cd+ab), a2+d2-b2-c2]])
def ApplyRotationMatrix(vector, rotationmatrix):
"""
This function take the output from the RotationMatrix function and
uses that to apply the rotation to an input vector
"""
a1 = (vector[0] * rotationmatrix[0, 0]) + (vector[1] * rotationmatrix[0, 1]) + (vector[2] * rotationmatrix[0, 2])
b1 = (vector[0] * rotationmatrix[1, 0]) + (vector[1] * rotationmatrix[1, 1]) + (vector[2] * rotationmatrix[1, 2])
c1 = (vector[0] * rotationmatrix[2, 0]) + (vector[1] * rotationmatrix[2, 1]) + (vector[2] * rotationmatrix[2, 2])
return np.array((a1, b1, c1)
'''
Functions for Calculating the angles of 3D vectors relative to one another
'''
def CalculateAngleBetweenVector(vector, vector2):
"""
Does what it says on the tin, outputs an angle in degrees between two input vectors.
"""
dp = np.dot(vector, vector2)
maga = math.sqrt((vector[0] ** 2) + (vector[1] ** 2) + (vector[2] ** 2))
magb = math.sqrt((vector2[0] ** 2) + (vector2[1] ** 2) + (vector2[2] ** 2))
magc = maga * magb
dpmag = dp / magc
#These if statements deal with rounding errors of floating point operations
if dpmag > 1:
error = dpmag - 1
print('error = {}, do not worry if this number is very small'.format(error))
dpmag = 1
elif dpmag < -1:
error = 1 + dpmag
print('error = {}, do not worry if this number is very small'.format(error))
dpmag = -1
angleindeg = ((math.acos(dpmag)) * 180) / math.pi
return angleindeg
def CalculateAngleAroundZ(Vector):
X,Y,Z = Vector[0], Vector[1], Vector[2]
AngleAroundZ = math.atan2(Y, X)
AngleAroundZdeg = (AngleAroundZ*180)/math.pi
return AngleAroundZdeg
def CalculateAngleAroundX(Vector):
X,Y,Z = Vector[0], Vector[1], Vector[2]
AngleAroundZ = math.atan2(Y, Z)
AngleAroundZdeg = (AngleAroundZ*180)/math.pi
return AngleAroundZdeg
def CalculateAngleAroundY(Vector):
X,Y,Z = Vector[0], Vector[1], Vector[2]
AngleAroundZ = math.atan2(X, Z)
AngleAroundZdeg = (AngleAroundZ*180)/math.pi
return AngleAroundZdeg
V1 = (0,0,1)
V2 = (3,5,4)
Xoffset = (CalculateAngleAroundX(V2))
Yoffset = (CalculateAngleAroundY(V2))
Zoffset = (CalculateAngleAroundZ(V2))
XRM = RotationMatrix((1,0,0), (Xoffset * 1))
YRM = RotationMatrix((0,1,0), (Yoffset * 1))
ZRM = RotationMatrix((0,0,1), (Zoffset * 1))
V2 = V2 / np.linalg.norm(V2)
V2X = ApplyRotationMatrix(V2, XRM)
V2XY = ApplyRotationMatrix(V2X, YRM)
V2XYZ = ApplyRotationMatrix(V2XY, ZRM)
print(V2XYZ)
print(CalculateAngleBetweenVector(V1, V2XYZ))
Any advice to fix this problem will be much appreciated.

I'm not sure to fully understand what you need but if it is to compute the angle between two vectors in space you can use the formula:
where a.b is the scalar product and theta is the angle between vectors.
thus your function CalculateAngleBetweenVector becomes:
def CalculateAngleBetweenVector(vector, vector2):
return math.acos(np.dot(vector,vector2)/(np.linalg.norm(vector)* np.linalg.norm(vector2))) * 180 /math.pi
You can also simplify your ApplyRotationMatrix function:
def ApplyRotationMatrix(vector, rotationmatrix):
"""
This function take the output from the RotationMatrix function and
uses that to apply the rotation to an input vector
"""
return rotationmatrix # vector
the # symbol is the matrix product
Hope this will help you. Feel free to precise your request if this is not helpfull.

Im an idiot I just needed to do the cross product and the dot product and rotate by the dot product *-1 around the cross product.

Electric force between particles using numpy arrays

I am trying to simulate a particle flying at another particle while undergoing electrical repulsion (or attraction), called Rutherford-scattering. I have succeeded in simulating (a few) particles using for loops and python lists. However, now I want to use numpy arrays instead. The model will use the following steps:
For all particles:
Calculate radial distance with all other particles
Calculate the angle with all other particles
Calculate netto force in x-direction and y-direction
Create matrix with netto xForce and yForce for each particle
Create accelaration (also x and y component) matrix by a = F/mass
Update speed matrix
Update position matrix
My problem is that I do not know how I can use numpy arrays in calculating the force components.
Here follows my code which is not runnable.
import numpy as np
# I used this function to calculate the force while using for-loops.
def force(x1, y1, x2, x2):
angle = math.atan((y2 - y1)/(x2 - x1))
dr = ((x1-x2)**2 + (y1-y2)**2)**0.5
force = charge2 * charge2 / dr**2
xforce = math.cos(angle) * force
yforce = math.sin(angle) * force
# The direction of force depends on relative location
if x1 > x2 and y1<y2:
xforce = xforce
yforce = yforce
elif x1< x2 and y1< y2:
xforce = -1 * xforce
yforce = -1 * yforce
elif x1 > x2 and y1 > y2:
xforce = xforce
yforce = yforce
else:
xforce = -1 * xforce
yforce = -1* yforce
return xforce, yforce
def update(array):
# this for loop defeats the entire use of numpy arrays
for particle in range(len(array[0])):
# find distance of all particles pov from 1 particle
# find all x-forces and y-forces on that particle
xforce = # sum of all x-forces from all particles
yforce = # sum of all y-forces from all particles
force_arr[0, particle] = xforce
force_arr[1, particle] = yforce
return force
# begin parameters
t = 0
N = 3
masses = np.ones(N)
charges = np.ones(N)
loc_arr = np.random.rand(2, N)
speed_arr = np.random.rand(2, N)
acc_arr = np.random.rand(2, N)
force = np.random.rand(2, N)
while t < 0.5:
force_arr = update(loc_arry)
acc_arr = force_arr / masses
speed_arr += acc_array
loc_arr += speed_arr
t += dt
# plot animation

One approach to model this problem with arrays may be:
define the point coordinates as a Nx2 array. (This will help with extensibility if you advance to 3-D points later)
define the intermediate variables distance, angle, force as NxN arrays to represent the pairwise interactions
Numpy things to know about:
You can call most numeric functions on arrays if the arrays have the same shape (or conforming shapes, which is a nontrivial topic...)
meshgrid helps you generate the array indices necessary to shapeshift your Nx2 arrays to compute NxN results
and a tangential note (ha ha) arctan2() computes a signed angle, so you can bypass the complex "which quadrant" logic
For example you can do something like this. Note in get_dist and get_angle the arithmetic operations between points take place in the bottom-most dimension:
import numpy as np
# 2-D locations of particles
points = np.array([[1,0],[2,1],[2,2]])
N = len(points) # 3
def get_dist(p1, p2):
r = p2 - p1
return np.sqrt(np.sum(r*r, axis=2))
def get_angle(p1, p2):
r = p2 - p1
return np.arctan2(r[:,:,1], r[:,:,0])
ii = np.arange(N)
ix, iy = np.meshgrid(ii, ii)
dist = get_dist(points[ix], points[iy])
angle = get_angle(points[ix], points[iy])
# ... compute force
# ... apply the force, etc.
For the sample 3-point vector shown above:
In [246]: dist
Out[246]:
array([[0. , 1.41421356, 2.23606798],
[1.41421356, 0. , 1. ],
[2.23606798, 1. , 0. ]])
In [247]: angle / np.pi # divide by Pi to make the numbers recognizable
Out[247]:
array([[ 0. , -0.75 , -0.64758362],
[ 0.25 , 0. , -0.5 ],
[ 0.35241638, 0.5 , 0. ]])

Here is one go with only a loop for each time step, and it should work for any number of dimensions, I have tested with 3 too:
from matplotlib import pyplot as plt
import numpy as np
fig, ax = plt.subplots()
N = 4
ndim = 2
masses = np.ones(N)
charges = np.array([-1, 1, -1, 1]) * 2
# loc_arr = np.random.rand(N, ndim)
loc_arr = np.array(((-1,0), (1,0), (0,-1), (0,1)), dtype=float)
speed_arr = np.zeros((N, ndim))
# compute charge matrix, ie c1 * c2
charge_matrix = -1 * np.outer(charges, charges)
time = np.linspace(0, 0.5)
dt = np.ediff1d(time).mean()
for i, t in enumerate(time):
# get (dx, dy) for every point
delta = (loc_arr.T[..., np.newaxis] - loc_arr.T[:, np.newaxis]).T
# calculate Euclidean distance
distances = np.linalg.norm(delta, axis=-1)
# and normalised unit vector
unit_vector = (delta.T / distances).T
unit_vector[np.isnan(unit_vector)] = 0 # replace NaN values with 0
# calculate force
force = charge_matrix / distances**2 # norm gives length of delta vector
force[np.isinf(force)] = 0 # NaN forces are 0
# calculate acceleration in all dimensions
acc = (unit_vector.T * force / masses).T.sum(axis=1)
# v = a * dt
speed_arr += acc * dt
# increment position, xyz = v * dt
loc_arr += speed_arr * dt
# plotting
if not i:
color = 'k'
zorder = 3
ms = 3
for i, pt in enumerate(loc_arr):
ax.text(*pt + 0.1, s='{}q {}m'.format(charges[i], masses[i]))
elif i == len(time)-1:
color = 'b'
zroder = 3
ms = 3
else:
color = 'r'
zorder = 1
ms = 1
ax.plot(loc_arr[:,0], loc_arr[:,1], '.', color=color, ms=ms, zorder=zorder)
ax.set_aspect('equal')
The above example produces, where the black and blue points signify the start and end positions, respectively:
And when charges are equal charges = np.ones(N) * 2 the system symmetry is preserved and the charges repel:
And finally with some random initial velocities speed_arr = np.random.rand(N, 2):
EDIT
Made a small change to the code above to make sure it was correct. (I was missing -1 on the resultant force, ie. force between +/+ should be negative, and I was summing down the wrong axis, apologies for that. Now in the cases where masses[0] = 5, the system evolves correctly:

The classic approach is to calculate electric field for all particles in the system. Say you have 3 charged particles all with positive charge:
particles = np.array([[1,0,0],[2,1,0],[2,2,0]]) # location of each particle
q = np.array([1,1,1]) # charge of each particle
The easiest way to compute the electric field at each particle`s location is for loop:
def for_method(pos,q):
"""Computes electric field vectors for all particles using for-loop."""
Evect = np.zeros( (len(pos),len(pos[0])) ) # define output electric field vector
k = 1 / (4 * np.pi * const.epsilon_0) * np.ones((len(pos),len(pos[0]))) * 1.602e-19 # make this into matrix as matrix addition is faster
# alternatively you can get rid of np.ones and just define this as a number
for i, v0 in enumerate(pos): # s_p - selected particle | iterate over all particles | v0 reference particle
for v, qc in zip(pos,q): # loop over all particles and calculate electric force sum | v particle being calculated for
if all((v0 == v)): # do not compute for the same particle
continue
else:
r = v0 - v #
Evect[i] += r / np.linalg.norm(r) ** 3 * qc #! multiply by charge
return Evect * k
# to find electric field at each particle`s location call
for_method(particles, q)
This function returns array of vectors with the same shape as input particles array. To find force on each, you simply multiply this vector with q array of charges. From there on, you can easily find your acceleration and integrate the system using your favourite ODE solver.
Performance Optimization & Accuracy
For method is the slowest possible approach. The field can be computed using solely linear algebra granting significant speed boost. Following code is very efficient Numpy matrix "one-liner" (almost one-liner) to this problem:
def CPU_matrix_method(pos,q):
"""Classic vectorization of for Coulomb law using numpy arrays."""
k = 1 / (4 * np.pi * const.epsilon_0) * np.ones((len(pos),3)) * 1.602e-19 # define electric constant
dist = distance.cdist(pos,pos) # compute distances
return k * np.sum( (( np.tile(pos,len(pos)).reshape((len(pos),len(pos),3)) - np.tile(pos,(len(pos),1,1))) * q.reshape(len(q),1)).T * np.power(dist,-3, where = dist != 0),axis = 1).T
Note that this and following code also return electric field vector for each particle.
You can get even higher performance if you offload this onto the GPU using Cupy library. Following code is almost identical to the CPU_matrix_method, I only expanded the one-liner a little so that you could see better what is going on:
def GPU_matrix_method(pos,q):
"""GPU Coulomb law vectorization.
Takes in numpy arrays, performs computations and returns cupy array"""
# compute distance matrix between each particle
k_cp = 1 / (4 * cp.pi * const.epsilon_0) * cp.ones((len(pos),3)) * 1.602e-19 # define electric constant, runs faster if this is matrix
dist = cp.array(distance.cdist(pos,pos)) # could speed this up with cupy cdist function! use this: cupyx.scipy.spatial.distance.cdist
pos, q = cp.array(pos), cp.array(q) # load inputs to GPU memory
dist_mod = cp.power(dist,-3) # compute inverse cube of distance
dist_mod[dist_mod == cp.inf] = 0 # set all infinity entries to 0 (i.e. diagonal elements/ same particle-particle pairs)
# compute by magic
return k_cp * cp.sum((( cp.tile(pos,len(pos)).reshape((len(pos),len(pos),3)) - cp.tile(pos,(len(pos),1,1))) * q.reshape(len(q),1)).T * dist_mod, axis = 1).T
Regarding the accuracy of the mentioned algorithms, if you compute the 3 methods on the particles array you get identical results:
[[-6.37828367e-10 -7.66608512e-10 0.00000000e+00]
[ 5.09048221e-10 -9.30757576e-10 0.00000000e+00]
[ 1.28780145e-10 1.69736609e-09 0.00000000e+00]]
Regarding the performance, I computed each algorithm on systems ranging from 2 to 5000 charged particles. Additionally I also included Numba precompiled version of the for_method to make the for-loop approach competitive:
We see that for-loop performs terribly needing over 400 seconds to compute for system with 5000 particles. Zooming in to the bottom part:
This shows that matrix approach to this problem is orders of magnitude better. To be exact the 5000 particle evaluation took 18.5s for Numba for-loop, 4s for CPU matrix(5 times faster than Numba), and 0.8s for GPU matrix* (23 times faster than Numba). The significant difference shows for larger arrays.
* GPU used was Nvidia K100.

Calculating mean value of a 2D array as a function of distance from the center in Python

I'm trying to calculate the mean value of a quantity(in the form of a 2D array) as a function of its distance from the center of a 2D grid. I understand that the idea is that I identify all the array elements that are at a distance R from the center, and then add them up and divide by the number of elements. However, I'm having trouble actually identifying an algorithm to go about doing this.
I have attached a working example of the code to generate the 2d array below. The code is for calculating some quantities that are resultant from gravitational lensing, so the way the array is made is irrelevant to this problem, but I have attached the entire code so that you could create the output array for testing.
import numpy as np
import multiprocessing
import matplotlib.pyplot as plt
n = 100 # grid size
c = 3e8
G = 6.67e-11
M_sun = 1.989e30
pc = 3.086e16 # parsec
Dds = 625e6*pc
Ds = 1726e6*pc #z=2
Dd = 1651e6*pc #z=1
FOV_arcsec = 0.0001
FOV_arcmin = FOV_arcsec/60.
pix2rad = ((FOV_arcmin/60.)/float(n))*np.pi/180.
rad2pix = 1./pix2rad
Renorm = (4*G*M_sun/c**2)*(Dds/(Dd*Ds))
#stretch = [10, 2]
# To create a random distribution of points
def randdist(PDF, x, n):
#Create a distribution following PDF(x). PDF and x
#must be of the same length. n is the number of samples
fp = np.random.rand(n,)
CDF = np.cumsum(PDF)
return np.interp(fp, CDF, x)
def get_alpha(args):
zeta_list_part, M_list_part, X, Y = args
alpha_x = 0
alpha_y = 0
for key in range(len(M_list_part)):
z_m_z_x = (X - zeta_list_part[key][0])*pix2rad
z_m_z_y = (Y - zeta_list_part[key][1])*pix2rad
alpha_x += M_list_part[key] * z_m_z_x / (z_m_z_x**2 + z_m_z_y**2)
alpha_y += M_list_part[key] * z_m_z_y / (z_m_z_x**2 + z_m_z_y**2)
return (alpha_x, alpha_y)
if __name__ == '__main__':
# number of processes, scale accordingly
num_processes = 1 # Number of CPUs to be used
pool = multiprocessing.Pool(processes=num_processes)
num = 100 # The number of points/microlenses
r = np.linspace(-n, n, n)
PDF = np.abs(1/r)
PDF = PDF/np.sum(PDF) # PDF should be normalized
R = randdist(PDF, r, num)
Theta = 2*np.pi*np.random.rand(num,)
x1= [R[k]*np.cos(Theta[k])*1 for k in range(num)]
y1 = [R[k]*np.sin(Theta[k])*1 for k in range(num)]
# Uniform distribution
#R = np.random.uniform(-n,n,num)
#x1= np.random.uniform(-n,n,num)
#y1 = np.random.uniform(-n,n,num)
zeta_list = np.column_stack((np.array(x1), np.array(y1))) # List of coordinates for the microlenses
x = np.linspace(-n,n,n)
y = np.linspace(-n,n,n)
X, Y = np.meshgrid(x,y)
M_list = np.array([0.1 for i in range(num)])
# split zeta_list, M_list, X, and Y
zeta_list_split = np.array_split(zeta_list, num_processes, axis=0)
M_list_split = np.array_split(M_list, num_processes)
X_list = [X for e in range(num_processes)]
Y_list = [Y for e in range(num_processes)]
alpha_list = pool.map(
get_alpha, zip(zeta_list_split, M_list_split, X_list, Y_list))
alpha_x = 0
alpha_y = 0
for e in alpha_list:
alpha_x += e[0]
alpha_y += e[1]
alpha_x_y = 0
alpha_x_x = 0
alpha_y_y = 0
alpha_y_x = 0
alpha_x_y, alpha_x_x = np.gradient(alpha_x*rad2pix*Renorm,edge_order=2)
alpha_y_y, alpha_y_x = np.gradient(alpha_y*rad2pix*Renorm,edge_order=2)
det_A = 1 - alpha_y_y - alpha_x_x + (alpha_x_x)*(alpha_y_y) - (alpha_x_y)*(alpha_y_x)
abs = np.absolute(det_A)
I = abs**(-1.)
O = np.log10(I+1)
plt.contourf(X,Y,O,100)
The array of interest is O, and I have attached a plot of how it should look like. It can be different based on the random distribution of points.
What I'm trying to do is to plot the mean values of O as a function of radius from the center of the grid. In the end, I want to be able to plot the average O as a function of distance from center in a 2d line graph. So I suppose the first step is to define circles of radius R, based on X and Y.
def circle(x,y):
r = np.sqrt(x**2 + y**2)
return r
Now I just have to figure out a way to find all the values of O, that have the same indices as equivalent values of R. Kinda confused on this part and would appreciate any help.

You can find the geometric coordinates of a circle with center (0,0) and radius R as such:
phi = np.linspace(0, 1, 50)
x = R*np.cos(2*np.pi*phi)
y = R*np.sin(2*np.pi*phi)
these values however will not fall on the regular pixel grid but in between.
In order to use them as sampling points you can either round the values and use them as indexes or interpolate the values from the near pixels.
Attention: The pixel indexes and the x, y are not the same. In your example (0,0) is at the picture location (50,50).

How can I vectorize this numpy rotation code?

The Goal:
I would like to vectorize (or otherwise speed up) this code. It rotates a 3d numpy model around its center point (let x,y,z denote the dimensions; then we want to rotate around the z-axis). The np model is binary voxels that are either "on" or "off"
I bet some basic matrix operation could do it, like take a layer and apply the rotation matrix to each element. The only issue with that is decimals; where should I have the new value land since cos(pi / 6) == sqrt(3) / 2?
The Code:
def rotate_model(m, theta):
'''
theta in degrees
'''
n =np.zeros(m.shape)
for i,layer in enumerate(m):
rotated = rotate(layer,theta)
n[i] = rotated
return n
where rotate() is:
def rotate(arr, theta):
'''
Rotates theta clockwise
rotated.shape == arr.shape, unlike scipy.ndimage.rotate(), which inflates size and also does some strange mixing
'''
if theta == int(theta):
theta *= pi / 180
theta = -theta
# theta=-theta b/c clockwise. Otherwise would default to counterclockwise
rotated =np.zeros(arr.shape)
#print rotated.shape[0], rotated.shape[1]
y_mid = arr.shape[0]//2
x_mid = arr.shape[1]//2
val = 0
for x_new in range(rotated.shape[1]):
for y_new in range(rotated.shape[0]):
x_centered = x_new - x_mid
y_centered = y_new - y_mid
x = x_centered*cos(theta) - y_centered*sin(theta)
y = x_centered*sin(theta) + y_centered*cos(theta)
x += x_mid
y += y_mid
x = int(round(x)); y = int(round(y)) # cast so range() picks it up
# lossy rotation
if x in range(arr.shape[1]) and y in range(arr.shape[0]):
val = arr[y,x]
rotated[y_new,x_new] = val
#print val
#print x,y
return rotated

You have a couple of problems in your code. First, if you want to fit the original image onto a rotated grid then you need a larger grid (usually). Alternatively, imagine a regular grid but the shape of your object - a rectangle - is rotated, thus becoming a "rhomb". It is obvious if you want to fit the entire rhomb - you need a larger output grid (array). On the other hand, you say in the code "rotated.shape == arr.shape, unlike scipy.ndimage.rotate(), which inflates size". If that is the case, maybe you do not want to fit the entire object? So, maybe it is OK to do this: rotated=np.zeros(arr.shape). But in general, yeah, one has to have a larger grid in order to fit the entire input image after it is rotated.
Another issue is angle conversion that you are doing:
if theta == int(theta):
theta *= pi / 180
theta = -theta
Why??? What will happen when I want to rotate the image by 1 radian? Or 2 radians? Am I forbidden to use integer number of radians? I think you are trying to do too much in this function and therefore it will be very confusing to do use it. Just require the caller to convert angles to radians. Or, you can do it inside this function if input theta is always in degrees. Or, you can add another parameter called, e.g., units and caller could set it to radians or degrees. Don't try to guess it based on "integer-ness" of input!
Now, let's rewrite your code a little bit:
rotated = np.zeros_like(arr) # instead of np.zero(arr.shape)
y_mid = arr.shape[0] // 2
x_mid = arr.shape[1] // 2
# val = 0 <- this is unnecessary
# pre-compute cos(theta) and sin(theta):
cs = cos(theta)
sn = sin(theta)
for x_new in range(rotated.shape[1]):
for y_new in range(rotated.shape[0]):
x = int(round((x_new - x_mid) * cs - (y_new - y_mid) * sn + x_mid)
y = int(round((x_new - x_mid) * sn - (y_new - y_mid) * cs + y_mid)
# just use comparisons, don't search through many values!
if 0 <= x < arr.shape[1] and 0 <= y < arr.shape[0]:
rotated[y_new, x_new] = arr[y, x]
So, now I can see (more easily) that for each pixel from the output array is mapped to a location in the input array. Yes, you can vectorize this.
import numpy as np
def rotate(arr, theta, unit='rad'):
# deal with theta units:
if unit.startswith('deg'):
theta = np.deg2rad(theta)
# for convenience, store array size:
ny, nx = arr.shape
# generate arrays of indices and flatten them:
y_new, x_new = np.indices(arr.shape)
x_new = x_new.ravel()
y_new = y_new.ravel()
# compute center of the array:
x0 = nx // 2
y0 = ny // 2
# compute old coordinates
xc = x_new - x0
yc = y_new - y0
x = np.round(np.cos(theta) * xc - np.sin(theta) * yc + x0).astype(np.int)
y = np.round(np.sin(theta) * xc - np.cos(theta) * yc + y0).astype(np.int)
# main idea to deal with indices is to create a mask:
mask = (x >= 0) & (x < nx) & (y >= 0) & (y < ny)
# ... and then select only those coordinates (both in
# input and "new" coordinates) that satisfy the above condition:
x = x[mask]
y = y[mask]
x_new = x_new[mask]
y_new = y_new[mask]
# map input values to output pixels *only* for selected "good" pixels:
rotated = np.zeros_like(arr)
rotated[y_new, x_new] = arr[y, x]
return rotated

Here is some code for anyone also doing 3d modeling. It solved my specific use-case pretty well. Still figuring out how to rotate in the proper plane. Hope it's helpful to you as well:
def rotate_model(m, theta):
'''
Redefines the prev 'rotate_model()' method
theta has to be in degrees
'''
rotated = scipy.ndimage.rotate(m, theta, axes=(1,2))
# have tried (1,0), (2,0), and now (1,2)
# ^ z is "up" and "2"
# scipy.ndimage.rotate() shrinks the model
# TODO: regrow it back
x_r = rotated.shape[1]
y_r = rotated.shape[0]
x_m = m.shape[1]
y_m = m.shape[0]
x_diff = abs(x_r - x_m)
y_diff = abs(y_r - y_m)
if x_diff%2==0 and y_diff%2==0:
return rotated[
x_diff//2 : x_r-x_diff//2,
y_diff//2 : y_r-y_diff//2,
:
]
elif x_diff%2==0 and y_diff%2==1:
# if this shift ends up turning the model to shit in a few iterations,
# change the following lines to include a flag that alternates cutting off the top and bottom bits of the array
return rotated[
x_diff//2 : x_r-x_diff//2,
y_diff//2+1 : y_r-y_diff//2,
:
]
elif x_diff%2==1 and y_diff%2==0:
return rotated[
x_diff//2+1 : x_r-x_diff//2,
y_diff//2 : y_r-y_diff//2,
:
]
else:
# x_diff%2==1 and y_diff%2==1:
return rotated[
x_diff//2+1 : x_r-x_diff//2,
y_diff//2+1 : y_r-y_diff//2,
:
]