diff --git a/exptool/analysis/pattern.py b/exptool/analysis/pattern.py index ebf0025..ae76e20 100644 --- a/exptool/analysis/pattern.py +++ b/exptool/analysis/pattern.py @@ -1,19 +1,11 @@ - -# 08-29-16: added maximum radius capabilities to bar_fourier_compute - -# 10-25-16: some redundancies noticed (bar_fourier_compute) and should be unified - ''' - pattern.py (part of exptool) tools to find patterns in the global simulation outputs - - - - +MSP 29 Aug 2016 Added maximum radius capabilities to bar_fourier_compute +MSP 25 Oct 2016 Some redundancies noticed (bar_fourier_compute) and should be unified BarTransform @@ -24,8 +16,10 @@ -Filtering algorithms for bar determination (e.g. look at better time-series algorithms) -Need a partial pattern calculator for bars that grow and disappear. Perhaps also in eof.py? +-Filter bad values from the pattern speed of the bar somehow +-Combine multiple coefficient series for a better estimate -BASIC USAGE: +BASIC USAGE Examples: # to transform a PSP output to have the bar on the X axis PSPTransform = pattern.BarTransform(PSPInput) @@ -37,10 +31,6 @@ ''' -from __future__ import absolute_import, division, print_function, unicode_literals - - - # general imports import time import numpy as np @@ -51,128 +41,242 @@ # exptool imports -from exptool.io import particle -from exptool.utils import kmeans -from exptool.utils import utils - - - - - - - -class BarTransform(): - ''' - BarTransform : class to do the work to calculate the bar position and transform particles - - on it's own, BarTransform will reset the particles to be in the bar frame (planar transformation) - - inputs - ----------------------- - ParticleInstanceIn : the input PSP instance - bar_angle : (default=None) the known bar angle - rel_bar_angle : (default=0.) the desired rotation angle relative to the bar major axis, counterclockwise (known or computed) - minr : (default=0.) the MINIMUM radius of particles to use to compute the bar angle - maxr : (default=1.) the MAXIMUM radius of particles to use in compute the bar angle - - outputs - ----------------------- - None - (ParticleInstanceIn will be modified to be in the planar bar transformation) - - - helper routines - ----------------------- - calculate_transform_and_return : overwrite the input PSP instance to have the raw positions be transformed - bar_fourier_compute : use m=2 fourier to transform to the bar frame +from ..io import particle + + +class BarFromCoefficients: + """ + Class to calculate and analyze the bar position from Fourier coefficients. + + Parameters + ---------- + times : array-like + Array of time values. + coefs : array-like + Array of complex Fourier coefficients. + unwrap_threshold : float, optional + Threshold for unwrapping the bar position phase (default is -1.). + smooth : bool, optional + Whether to smooth the unwrapped position (default is False). + reverse : bool, optional + Whether to reverse the unwrapping direction (default is False). + adjust : float, optional + Adjustment value used in unwrapping (default is np.pi). + verbose : int, optional + Verbosity level (default is 0). + smth_derivative : int, optional + Smoothing factor for the polynomial derivative calculation (default is 0). + spline_derivative : int, optional + Smoothing factor for the spline derivative calculation (default is 0). + """ + def __init__(self, times, coefs, unwrap_threshold=-np.pi/2., smooth=False, reverse=False, adjust=2*np.pi, verbose=0, smth_derivative=0, spline_derivative=0): + self.verbose = verbose + self.time = times + self.cos = np.real(coefs) + self.sin = np.imag(coefs) + self.barposition = np.arctan2(self.sin, self.cos) + # compute the unwrapped position + self._unwrap_position(unwrap_threshold, smooth, reverse, adjust) + # compute the derivative + self._frequency_and_derivative(smth_derivative,spline_derivative) + + def _unwrap_position(self, unwrap_threshold, smooth, reverse, adjust): + """ + Unwrap the bar position to avoid discontinuities. + + Parameters + ---------- + unwrap_threshold : float + Threshold for phase unwrapping. + smooth : bool + Whether to smooth the unwrapped position. + reverse : bool + Whether to reverse the unwrapping direction. + adjust : float + Adjustment value used in unwrapping. + """ + running_number_of_rotations = 0 + number_of_rotations = np.zeros_like(self.barposition) + + # which way are we rotating + primarydirection = np.nanmedian(np.ediff1d(self.barposition)) + + # start from the beginning and keep track of number of rotations + for i in range(1, len(self.barposition)): + if (primarydirection > 0): + if (self.barposition[i] - self.barposition[i-1]) < unwrap_threshold: + running_number_of_rotations += 1 + else: + if (self.barposition[i] - self.barposition[i-1]) > -1. * unwrap_threshold: + running_number_of_rotations += 1 + number_of_rotations[i] = running_number_of_rotations + + # now straighten out the zeros + if (primarydirection > 0): + unwrapped_barposition = self.barposition + number_of_rotations * adjust + else: + unwrapped_barposition = - self.barposition + number_of_rotations * adjust + self.pos = unwrapped_barposition + + def _frequency_and_derivative(self, smth_derivative,spline_derivative): + """ + Calculate the frequency and derivative of the unwrapped bar position. + + Parameters + ---------- + spline_derivative : int + Smoothing factor for the spline derivative calculation. + """ + # make a numerical derivative estimate + self.deriv = np.zeros_like(self.pos) + for i in range(1, len(self.pos) - 1): + self.deriv[i] = (self.pos[i+1] - self.pos[i-1]) / (2 * (self.time[i] - self.time[i-1])) + + if (smth_derivative): + smth_params = np.polyfit(self.time, self.deriv, smth_derivative) + pos_func = np.poly1d(smth_params) + self.deriv = pos_func(self.time) - ''' + # hard set as a cubic spline, + # number is a smoothing factor between knots, see scipy.UnivariateSpline + # + # recommended: 7 for dt=0.002 spacing + if spline_derivative: + spl = UnivariateSpline(self.time, self.pos, k=3, s=spline_derivative) + self.deriv = (spl.derivative())(self.time) + self.dderiv = np.zeros_like(self.deriv) + # + # can also do a second deriv + for indx, timeval in enumerate(self.time): + self.dderiv[indx] = spl.derivatives(timeval)[2] + + def print_bar(self, outfile): + """ + Print the bar position, its derivative, and time to a file. + + Parameters + ---------- + outfile : str + Path to the output file. + """ + with open(outfile, 'w') as f: + for i in range(len(self.time)): + print(self.time[i], self.pos[i], self.deriv[i], file=f) + return None - def __init__(self,ParticleInstanceIn,bar_angle=None,rel_bar_angle=0.,minr=0.,maxr=1.): - ''' - see documentation above - ''' +class BarTransform: + """ + BarTransform: A class to calculate the bar position and transform particles into the bar frame. + + On its own, BarTransform will reset the particles to be in the bar frame (planar transformation). + + Parameters + ---------- + ParticleInstanceIn : object + The input particle instance. + bar_angle : float, optional + The known bar angle. If None, it will be computed (default is None). + rel_bar_angle : float, optional + The desired rotation angle relative to the bar major axis, counterclockwise (default is 0.). + minr : float, optional + The minimum radius of particles to use to compute the bar angle (default is 0.). + maxr : float, optional + The maximum radius of particles to use to compute the bar angle (default is 1.). + + Attributes + ---------- + ParticleInstanceIn : object + The input particle instance. + bar_angle : float + The computed or provided bar angle. + data : dict + Dictionary containing the transformed particle data. + time : float + The time of the particle instance. + filename : str + The filename of the particle instance. + comp : str + The component of the particle instance. + + Methods + ------- + calculate_transform_and_return() + Modify the input particle instance to be in the bar frame. + bar_fourier_compute(posx, posy, minr=0., maxr=1.) + Use m=2 Fourier analysis to compute the bar angle. + """ + + def __init__(self, ParticleInstanceIn, bar_angle=None, rel_bar_angle=0., minr=0., maxr=1.): self.ParticleInstanceIn = ParticleInstanceIn - self.bar_angle = bar_angle - self.data = dict() - if self.bar_angle == None: - self.bar_angle = -1.*BarTransform.bar_fourier_compute(self,self.ParticleInstanceIn.data['x'],self.ParticleInstanceIn.data['y'],maxr=maxr) - - #-1.*BarTransform.bar_fourier_compute(self,self.ParticleInstanceIn.xpos,self.ParticleInstanceIn.ypos,maxr=maxr) + if self.bar_angle is None: + self.bar_angle = -1. * self.bar_fourier_compute(self.ParticleInstanceIn.data['x'], self.ParticleInstanceIn.data['y'], minr=minr, maxr=maxr) - # do an arbitary rotation of the particles relative to the bar? + # Apply relative bar angle rotation self.bar_angle += rel_bar_angle + # Perform the transformation self.calculate_transform_and_return() - def calculate_transform_and_return(self): - ''' - calculate_transform_and_return - do the modification of the input PSP instance to be in the bar frame. - - inputs - ---------------------------- - self (BarTransform) - - - ''' - - - transformed_x = self.ParticleInstanceIn.data['x']*np.cos(self.bar_angle) - self.ParticleInstanceIn.data['y']*np.sin(self.bar_angle) - #self.ParticleInstanceIn.xpos*np.cos(self.bar_angle) - self.ParticleInstanceIn.ypos*np.sin(self.bar_angle) - transformed_y = self.ParticleInstanceIn.data['x']*np.sin(self.bar_angle) + self.ParticleInstanceIn.data['y']*np.cos(self.bar_angle) - #self.ParticleInstanceIn.xpos*np.sin(self.bar_angle) + self.ParticleInstanceIn.ypos*np.cos(self.bar_angle) - - transformed_vx = self.ParticleInstanceIn.data['vx']*np.cos(self.bar_angle) - self.ParticleInstanceIn.data['vy']*np.sin(self.bar_angle) - #self.ParticleInstanceIn.xvel*np.cos(self.bar_angle) - self.ParticleInstanceIn.yvel*np.sin(self.bar_angle) - transformed_vy = self.ParticleInstanceIn.data['vx']*np.sin(self.bar_angle) + self.ParticleInstanceIn.data['vy']*np.cos(self.bar_angle) - #self.ParticleInstanceIn.xvel*np.sin(self.bar_angle) + self.ParticleInstanceIn.yvel*np.cos(self.bar_angle) - - + """ + Modify the input particle instance to be in the bar frame. + """ + # Transform positions + transformed_x = self.ParticleInstanceIn.data['x'] * np.cos(self.bar_angle) - self.ParticleInstanceIn.data['y'] * np.sin(self.bar_angle) + transformed_y = self.ParticleInstanceIn.data['x'] * np.sin(self.bar_angle) + self.ParticleInstanceIn.data['y'] * np.cos(self.bar_angle) + + # Transform velocities + transformed_vx = self.ParticleInstanceIn.data['vx'] * np.cos(self.bar_angle) - self.ParticleInstanceIn.data['vy'] * np.sin(self.bar_angle) + transformed_vy = self.ParticleInstanceIn.data['vx'] * np.sin(self.bar_angle) + self.ParticleInstanceIn.data['vy'] * np.cos(self.bar_angle) + + # Update the data dictionary self.data['x'] = transformed_x self.data['y'] = transformed_y self.data['z'] = np.copy(self.ParticleInstanceIn.data['z']) - #np.copy(self.ParticleInstanceIn.zpos) # interesting. needs to be a copy for later operations to work! - self.data['vx'] = transformed_vx self.data['vy'] = transformed_vy self.data['vz'] = np.copy(self.ParticleInstanceIn.data['vz']) - #np.copy(self.ParticleInstanceIn.zvel) - self.data['m'] = self.ParticleInstanceIn.data['m'] - #self.ParticleInstanceIn.mass self.data['potE'] = self.ParticleInstanceIn.data['potE'] - #self.ParticleInstanceIn.pote + # Update metadata self.time = self.ParticleInstanceIn.time self.filename = self.ParticleInstanceIn.filename self.comp = self.ParticleInstanceIn.comp + def bar_fourier_compute(self, posx, posy, minr=0., maxr=1.): + """ + Use x and y positions to compute the m=2 Fourier phase angle. - def bar_fourier_compute(self,posx,posy,minr=0.,maxr=1.): - ''' - - use x and y positions to compute the m=2 power, and find phase angle + Parameters + ---------- + posx : array-like + x positions of particles. + posy : array-like + y positions of particles. + minr : float, optional + Minimum radius to consider (default is 0.). + maxr : float, optional + Maximum radius to consider (default is 1.). - TODO: - generalize to transform to any azimuthal order? - - ''' - w = np.where( ( (posx*posx + posy*posy)**0.5 > minr ) & ((posx*posx + posy*posy)**0.5 < maxr ))[0] - - aval = np.sum( np.cos( 2.*np.arctan2(posy[w],posx[w]) ) ) - bval = np.sum( np.sin( 2.*np.arctan2(posy[w],posx[w]) ) ) - - return np.arctan2(bval,aval)/2. + Returns + ------- + float + The m=2 phase angle. + """ + radius = np.sqrt(posx**2 + posy**2) + w = np.where((radius > minr) & (radius < maxr))[0] + aval = np.sum(np.cos(2. * np.arctan2(posy[w], posx[w]))) + bval = np.sum(np.sin(2. * np.arctan2(posy[w], posx[w]))) + return np.arctan2(bval, aval) / 2. @@ -191,9 +295,9 @@ def __init__(self,**kwargs): try: # check to see if bar file has already been created self.read_bar(kwargs['file']) - print('pattern.BarDetermine: BarInstance sucessfully read.') + print('pattern.BarDetermine: BarInstance successfully read.') - except: + except (FileNotFoundError, IOError, ValueError): print('pattern.BarDetermine: no compatible bar file found.') @@ -453,7 +557,7 @@ def read_bar(self,infile): pos.append(q[1]) try: deriv.append(q[2]) - except: + except IndexError: pass self.time = np.array(time) @@ -468,188 +572,239 @@ def read_bar(self,infile): +def find_barangle(time, BarInstance, interpolate=True): + """ + Use a bar instance to match the output time to a bar position. + Parameters + ---------- + time : array-like + Array of time values at which to find the bar position. + BarInstance : object + An instance of a class (such as `BarFromCoefficients`) that contains + bar positions and corresponding times. + interpolate : bool, optional + Whether to interpolate the bar position using a spline. If False, + the function finds the closest available bar position (default is True). + Returns + ------- + indx_barpos : array-like + Array of bar positions corresponding to the input time values. -def compute_bar_lag(ParticleInstance,rcut=0.01,verbose=0): - ''' - # - # simple fourier method to calculate where the particles are in relation to the bar - # - ''' - R = (ParticleInstance.data['x']*ParticleInstance.data['x'] + ParticleInstance.data['y']*ParticleInstance.data['y'])**0.5 - #(ParticleInstance.xpos*ParticleInstance.xpos + ParticleInstance.ypos*ParticleInstance.ypos)**0.5 - TH = np.arctan2(ParticleInstance.data['y'],ParticleInstance.data['x']) - #np.arctan2(ParticleInstance.ypos,ParticleInstance.xpos) - loR = np.where( R < rcut)[0] - A2 = np.sum(ParticleInstance.mass[loR] * np.cos(2.*TH[loR])) - B2 = np.sum(ParticleInstance.mass[loR] * np.sin(2.*TH[loR])) - bar_angle = 0.5*np.arctan2(B2,A2) - - if (verbose): - print('Position angle is {0:4.3f} . . .'.format(bar_angle)) - - # - # two steps: - # 1. rotate theta so that the bar is aligned at 0,2pi - # 2. fold onto 0,pi to compute the lag - # - tTH = (TH - bar_angle + np.pi/2.) % np.pi # compute lag with bar at pi/2 - # - # verification plot - #plt.scatter( R[0:10000]*np.cos(tTH[0:10000]-np.pi/2.),R[0:10000]*np.sin(tTH[0:10000]-np.pi/2.),color='black',s=0.5) - return tTH - np.pi/2. # retransform to bar at 0 - + Notes + ----- + This function can take arrays as input. It currently handles only one + direction of bar position matching and places a guard against NaN values + in the bar positions. + """ + # Place a guard against NaN values + BarInstance.pos = np.nan_to_num(BarInstance.pos, nan=0.0) + sord = 0 # Should this be a variable? + if interpolate: + not_nan = np.where(~np.isnan(BarInstance.pos)) + bar_func = UnivariateSpline(BarInstance.time[not_nan], -BarInstance.pos[not_nan], s=sord) - -def find_barangle(time,BarInstance,interpolate=True): - ''' - # - # use a bar instance to match the output time to a bar position - # - # can take arrays! - # - # but feels like it only goes one direction? - # - ''' - # place in a guard against nan values - BarInstance.pos[BarInstance.pos == np.nan] = 0. - - - sord = 0 # should this be a variable? - # - if (interpolate): - not_nan = np.where(np.isnan(BarInstance.pos)==False) - bar_func = UnivariateSpline(BarInstance.time[not_nan],-BarInstance.pos[not_nan],s=sord) - # try: indx_barpos = np.zeros([len(time)]) - for indx,timeval in enumerate(time): - # - if (interpolate): + for indx, timeval in enumerate(time): + if interpolate: indx_barpos[indx] = bar_func(timeval) - # - # else: - indx_barpos[indx] = -BarInstance.pos[ abs(timeval-BarInstance.time).argmin()] - # - except: - if (interpolate): + indx_barpos[indx] = -BarInstance.pos[np.abs(timeval - BarInstance.time).argmin()] + except TypeError: + if interpolate: indx_barpos = bar_func(time) - # else: - indx_barpos = -BarInstance.pos[ abs(time-BarInstance.time).argmin()] - # - return indx_barpos - + indx_barpos = -BarInstance.pos[np.abs(time - BarInstance.time).argmin()] + return indx_barpos -def find_barpattern(intime,BarInstance,smth_order=2): - ''' - # - # use a bar instance to match the output time to a bar pattern speed - # - # simple differencing--may want to be careful with this. - # needs a guard for the end points - # - ''' - - # grab the derivative at whatever smoothing order - BarInstance.frequency_and_derivative(smth_order=smth_order) +def find_barpattern(intime, BarInstance, smth_order=2): + """ + Use a bar instance to match the output time to a bar pattern speed. + + Parameters + ---------- + intime : array-like + Array of time values at which to find the bar pattern speed. + BarInstance : object + An instance of a class (such as `BarFromCoefficients`) that contains + bar positions, times, and their derivatives. + smth_order : int, optional + Smoothing factor for the derivative calculation (default is 2). + + Returns + ------- + barpattern : array-like + Array of bar pattern speeds corresponding to the input time values. + """ + + # Compute the derivative of the bar position at the specified smoothing order + BarInstance._frequency_and_derivative(spline_derivative=smth_order) try: - + # Initialize an array to hold the bar pattern speeds barpattern = np.zeros([len(intime)]) - for indx,timeval in enumerate(intime): - - best_time = abs(timeval-BarInstance.time).argmin() + # Loop over each time value in the input array + for indx, timeval in enumerate(intime): + # Find the index of the closest time in BarInstance.time + best_time = abs(timeval - BarInstance.time).argmin() + # Get the derivative (bar pattern speed) at the closest time barpattern[indx] = BarInstance.deriv[best_time] - except: - - best_time = abs(intime-BarInstance.time).argmin() + except TypeError: + # Handle the case where intime is a single value + best_time = abs(intime - BarInstance.time).argmin() + # Get the derivative (bar pattern speed) at the closest time barpattern = BarInstance.deriv[best_time] return barpattern -'''Not sure if this is the best place for this - wrote code to make a barfile using fourier -analysis to find m=2 phase angle + then pattern speed based on this angle. This is to replace -the EOF info if the EOF info is weird. Output file formats should be identical''' -class fourier_barfiles(): - def parse_list(self): - - f = open(self.slist) - s_list = [] - for line in f: - d = [q for q in line.split()] - s_list.append(d[0]) +class BarFromFourier: + """ + Class to compute the bar pattern speed and phase angle using Fourier analysis. + + This class replaces the EOF information if it appears to be incorrect. The output file + formats are designed to be identical to those generated using EOF information. + + Parameters + ---------- + inputfiles : str + Path to a file containing a list of input files to be processed. + + Attributes + ---------- + slist : str + Path to the input files list. + SLIST : array-like + List of input files parsed from the file. + pos : array-like + Array of bar positions. + deriv : array-like + Array of bar pattern speeds (derivatives). + time : array-like + Array of time steps corresponding to the bar positions and speeds. + """ + + def __init__(self, inputfiles): + self.slist = inputfiles + def parse_list(self): + """ + Parse the list of input files from the provided file path. + + This method reads the file specified in `self.slist` and stores the list + of input files in the attribute `self.SLIST`. + """ + with open(self.slist, 'r') as f: + s_list = [line.split()[0] for line in f] self.SLIST = np.array(s_list) - - def bar_fourier_compute(self,posx,posy,maxr=0.5, minr=.001): - - # - # use x and y positions tom compute the m=2 power, and find phase angle - # - w = np.where( ((posx*posx + posy*posy)**0.5 < maxr) & - ((posx*posx + posy*posy)**0.5 > minr) )[0] - - aval = np.sum( np.cos( 2.*np.arctan2(posy[w],posx[w]) ) ) - bval = np.sum( np.sin( 2.*np.arctan2(posy[w],posx[w]) ) ) - - return np.arctan2(bval,aval)/2. + def bar_fourier_compute(self, posx, posy, maxr=0.5, minr=0.001): + """ + Compute the m=2 Fourier phase angle from particle positions. + + Parameters + ---------- + posx : array-like + x positions of particles. + posy : array-like + y positions of particles. + maxr : float, optional + Maximum radius to consider (default is 0.5). + minr : float, optional + Minimum radius to consider (default is 0.001). + + Returns + ------- + float + The m=2 phase angle. + """ + # Select particles within the specified radius range + radius = np.sqrt(posx**2 + posy**2) + w = np.where((radius < maxr) & (radius > minr))[0] + + # Compute m=2 Fourier components + aval = np.sum(np.cos(2. * np.arctan2(posy[w], posx[w]))) + bval = np.sum(np.sin(2. * np.arctan2(posy[w], posx[w]))) + + return np.arctan2(bval, aval) / 2. def bar_speed(self, filelist, comp='star'): + """ + Compute the bar pattern speed from the list of input files. + + Parameters + ---------- + filelist : array-like + List of input files to process. + comp : str, optional + Component to analyze (default is 'star'). + + Returns + ------- + dict + Dictionary containing arrays of time steps, bar positions, and bar pattern speeds. + """ self.slist = filelist - fourier_barfiles.parse_list(self) - pos = particle.Input(self.SLIST[0],comp=comp,verbose=0) - pos_p1 = particle.Input(self.SLIST[1],comp=comp,verbose=0) + self.parse_list() + + pos = particle.Input(self.SLIST[0], comp=comp, verbose=0) + pos_p1 = particle.Input(self.SLIST[1], comp=comp, verbose=0) first_bar_angle = self.bar_fourier_compute(pos.data['x'], pos.data['y']) - #get time step + + # Calculate the timestep timestep = pos_p1.time - pos.time tt = np.array([]) pp = np.array([]) rot = np.array([]) - for i in range(0,len(self.SLIST)): - #loop through snapshot files in simulation, open file - pos = particle.Input(self.SLIST[i],comp=comp,verbose=0) - #compute bar angle + + for i in range(len(self.SLIST)): + pos = particle.Input(self.SLIST[i], comp=comp, verbose=0) bar_angle = self.bar_fourier_compute(pos.data['x'], pos.data['y']) - #if first time step, old bar angle = current bar angle + if i == 0: old_bar_angle = first_bar_angle - pattern_speed = (old_bar_angle-bar_angle)/(timestep) - #bar_angle > old_bar_angle, if difference is near 180, flip it (took this bit from rachel) - if abs(old_bar_angle-bar_angle)>=(np.pi*3/4): - pattern_speed = (old_bar_angle - (bar_angle + np.pi))/(timestep) + + pattern_speed = (old_bar_angle - bar_angle) / timestep + + if abs(old_bar_angle - bar_angle) >= (np.pi * 3 / 4): + pattern_speed = (old_bar_angle - (bar_angle + np.pi)) / timestep + pp = np.append(pp, bar_angle) rot = np.append(rot, pattern_speed) tt = np.append(tt, pos.time) - #use this bar angle as 'old' angle for next step + old_bar_angle = bar_angle + self.pos = pp self.deriv = rot self.time = tt - return {'time':tt,'pos':pp, 'deriv':rot} - - def print_bar(self,simulation_directory,simulation_name): - - # - # print the barfile to file - # - - - f = open(simulation_directory+simulation_name+'fourier_barpos.dat','w') - for i in range(0,len(self.SLIST)): - print(self.time[i],self.pos[i],self.deriv[i],end="\n",file=f) - - f.close() - - return None \ No newline at end of file + + return {'time': tt, 'pos': pp, 'deriv': rot} + + def print_bar(self, simulation_directory, simulation_name): + """ + Print the bar positions and pattern speeds to a file. + + Parameters + ---------- + simulation_directory : str + Directory where the output file will be saved. + simulation_name : str + Base name for the output file. + """ + output_file = simulation_directory + simulation_name + 'fourier_barpos.dat' + + with open(output_file, 'w') as f: + for i in range(len(self.SLIST)): + print(self.time[i], self.pos[i], self.deriv[i], file=f) + + return None diff --git a/exptool/analysis/trapping.py b/exptool/analysis/trapping.py index 158f9dd..6ec3aad 100644 --- a/exptool/analysis/trapping.py +++ b/exptool/analysis/trapping.py @@ -5,6 +5,7 @@ MSP 23 Dec 2017 Break out bar finding algorithms to the more general pattern.py MSP 1 Mar 2019 Work on homogenizing docstrings and general commenting MSP 27 Oct 2021 Enable flexible particle number handling +MSP 18 May 2024 Create HDF5 input/ouput CLASSES: @@ -50,6 +51,9 @@ import os from scipy import interpolate +# io import +import h5py + # multiprocessing imports import itertools from multiprocessing import Pool, freeze_support @@ -109,7 +113,7 @@ def accept_files(self,filelist,verbose=0): ApsFinding.parse_list(self) - def parse_list(self): + def _parse_list(self): """ parse files from the input list @@ -128,7 +132,7 @@ def parse_list(self): self.SLIST = np.array(s_list) if self.verbose >= 1: - print('ApsFinding.parse_list: Accepted {0:d} files.'.format(len(self.SLIST))) + print('exptool.trapping.ApsFinding.parse_list: Accepted {0:d} files.'.format(len(self.SLIST))) def determine_r_aps(self,filelist,comp,particle_indx=-1,runtag='',out_directory='',threedee=False,return_aps=False,changingindx=False): @@ -141,21 +145,25 @@ def determine_r_aps(self,filelist,comp,particle_indx=-1,runtag='',out_directory= # take the inputs and identify all files that we will loop through self.slist = filelist - ApsFinding.parse_list(self) + ApsFinding._parse_list(self) # now we have self.SLIST, the parsed list of files we will analyse # first, check type of particle_index if isinstance(particle_indx,int): if particle_indx < 0: # if particle_indx < 0, make the comparison index all particles - Oa = particle.Input(self.SLIST[0],legacy=False,comp=comp,verbose=0) - particle_indx = np.arange(0,Oa.nbodies,1) + Oa = particle.Input(self.SLIST[0],comp=comp,verbose=0) + particle_indx = np.arange(0,Oa.data['id'].size,1) else: # limit to the maximum number desired particle_indx = np.arange(0,particle_indx,1) + + # assume an array has been passed + elif isinstance(particle_indx,np.ndarray): + changingindx = True + else: - # assume an array has been passed and accept: could check - pass + raise ValueError("exptool.ApsFinding.trapping._determine_r_aps: particle_indx must be an integer or an array.") # sort the particle indices particle_indx = particle_indx[particle_indx.argsort()] @@ -167,21 +175,25 @@ def determine_r_aps(self,filelist,comp,particle_indx=-1,runtag='',out_directory= # stamps the output file with the current time. do we like this? # tstamp = datetime.datetime.strftime(datetime.datetime.now(), '%Y-%m-%d+%H:%M:%S') + #tstamp = datetime.datetime.strftime(datetime.datetime.now(), '%Y-%m-%d+%H') # create a new file using the particle number, runtag, and time - outputfile = out_directory+'RadialAps_N{}_r{}_T{}.dat'.format(total_orbits,runtag,tstamp) - f = open(outputfile,'wb+') + outputfile = out_directory+'RadialAps_N{}_r{}_T{}.h5'.format(total_orbits,runtag,tstamp) + f = h5py.File(outputfile,"w") - # - # print descriptor string - # + # create descriptor string desc = 'apsfile for '+comp+' in '+out_directory+', norbits='+str(total_orbits)+', threedee='+str(threedee)+', using '+filelist - np.array([desc],dtype='S200').tofile(f) + + # Write the descriptor string as an attribute + f.attrs['description'] = desc - aps_dictionary = dict() # make blank dictionary for the aps - for i in range(0,total_orbits): aps_dictionary[i] = [] + # make blank dictionary for the aps + aps_dictionary = dict() + # make a blank array for each orbit + for i in particle_indx: aps_dictionary[i] = [] + # loop through files for i in range(1,len(self.SLIST)-1): if i==1: @@ -213,16 +225,16 @@ def determine_r_aps(self,filelist,comp,particle_indx=-1,runtag='',out_directory= X2 = Ob.data['x'];Y2 = Ob.data['y'];Z2 = Ob.data['z'];I2 = Ob.data['id'] X3 = Oc.data['x'];Y3 = Oc.data['y'];Z3 = Oc.data['z'];I3 = Oc.data['id'] + # compute radial positions if threedee: - R1 = np.sqrt(X1*X1 + Y1*Y1 + Z1*Z1) - R2 = np.sqrt(X2*X2 + Y2*Y2 + Z2*Z2) - R3 = np.sqrt(X3*X3 + Y3*Y3 + Z3*Z3) - + R1 = np.linalg.norm([X1,Y1,Z1],axis=0) + R2 = np.linalg.norm([X2,Y2,Z2],axis=0) + R3 = np.linalg.norm([X3,Y3,Z3],axis=0) else: - R1 = np.sqrt(X1*X1 + Y1*Y1) - R2 = np.sqrt(X2*X2 + Y2*Y2) - R3 = np.sqrt(X3*X3 + Y3*Y3) + R1 = np.linalg.norm([X1,Y1],axis=0) + R2 = np.linalg.norm([X2,Y2],axis=0) + R3 = np.linalg.norm([X3,Y3],axis=0) else: # i!=1 @@ -252,9 +264,9 @@ def determine_r_aps(self,filelist,comp,particle_indx=-1,runtag='',out_directory= X3 = Oc.data['x'];Y3 = Oc.data['y'];Z3 = Oc.data['z'];I3 = Oc.data['id'] if threedee: - R3 = np.sqrt(X3*X3 + Y3*Y3 + Z3*Z3) + R3 = np.linalg.norm([X3,Y3,Z3],axis=0) else: - R3 = np.sqrt(X3*X3 + Y3*Y3) + R3 = np.linalg.norm([X3,Y3],axis=0) # R1 might be the shortest, if particles are added, so only compare up to the length of r1 @@ -276,59 +288,61 @@ def determine_r_aps(self,filelist,comp,particle_indx=-1,runtag='',out_directory= #orderid = IN[:r1length][aps] if self.verbose > 0: - print('Current time: {4.3f}'.format(tval),end='\r', flush=True) - - # under this convention, the user needs to keep track of the particle index that was input - #for j in range(0,len(index_tags)): - # aps_dictionary[orderid[j]].append([tval,x[j],y[j],z[j]]) + print('exptool.ApsFinding.trapping._determine_r_aps: Current time: {0:4.3f}'.format(tval),end='\r', flush=True) - # under this convention, the id of the orbit is preserved and used as the dictionary key - for j in range(0,len(index_tags)): + # the id of the orbit is preserved and used as the dictionary key + for j in range(0,len(id)): aps_dictionary[id[j]].append([tval,x[j],y[j],z[j]]) - + # create a tracker for the number of aps per orbit self.napsides = np.zeros([total_orbits,2]) # print a header with the number of orbits - np.array([total_orbits],dtype='i').tofile(f) + f.attrs['total_orbits'] = total_orbits orbits_with_apocentre = 0 + # go back through all the orbits and write to file for j in range(0,total_orbits): orbit_aps_array = np.array(aps_dictionary[particle_indx[j]]) - # print the index to the file - np.array([particle_indx[j]],dtype='i').tofile(f) - # if there are valid turning points: if (len(orbit_aps_array) > 0): orbits_with_apocentre += 1 - naps = len(orbit_aps_array[:,0]) # this might be better as shape - np.array([naps],dtype='i').tofile(f) + # count the number of turning points + naps = len(orbit_aps_array[:,0]) - self.napsides[j,0] = naps - self.napsides[j,1] = len(orbit_aps_array.reshape(-1,)) + # create a dataset with the index of the particle as the tag + dataset = f.create_dataset(str(particle_indx[j]), data=orbit_aps_array) + + # create attributes for the dataset + dataset.attrs['naps'] = naps - np.array( orbit_aps_array.reshape(-1,),dtype='f').tofile(f) # no valid turning points: put in a blank else: + # create a dataset with the index of the particle as the tag + dataset = f.create_dataset(str(particle_indx[j]), data=np.array([-1.])) + + # create attributes for the dataset + dataset.attrs['naps'] = 0 + # guard against zero length - np.array([1],dtype='i').tofile(f) + #np.array([1],dtype='i').tofile(f) # indices start at 1 - np.array( np.array(([-1.,-1.,-1.,-1.])).reshape(-1,),dtype='f').tofile(f) + #np.array( np.array(([-1.,-1.,-1.,-1.])).reshape(-1,),dtype='f').tofile(f) f.close() - print('trapping.ApsFinding.determine_r_aps: found {} orbits (out of {}) with valid apocentres.'.format(orbits_with_apocentre,total_orbits)) + print('exptool.trapping.ApsFinding.determine_r_aps: found {} orbits (out of {}) with valid apocentres.'.format(orbits_with_apocentre,total_orbits)) - print('trapping.ApsFinding.determine_r_aps: savefile is {}'.format(outputfile)) + print('exptool.trapping.ApsFinding.determine_r_aps: savefile is {}'.format(outputfile)) if (return_aps): ApsDict = ApsFinding.read_aps_file(self,outputfile) @@ -426,324 +440,272 @@ def read_trapping_file(t_file,tdtype='i1'): -def reduce_aps_dictionary(TrappingInstance,norb): - ''' - sometimes you just don't need all those apsides - ''' +def reduce_aps_dictionary(TrappingInstance, norb): + """ + Reduces the apsides data in a trapping instance to only include a specified number of orbits. + + This function takes a trapping instance dictionary and reduces it to include only the + specified number of orbits (`norb`). It retains the description and the first `norb` + orbits' data. + + Parameters + ---------- + TrappingInstance : dict + A dictionary containing the trapping instance data with multiple orbits. + norb : int + The number of orbits to include in the reduced dictionary. + + Returns + ------- + TrappingInstanceOut : dict + A reduced dictionary containing only the specified number of orbits + from the original trapping instance. + + Examples + -------- + >>> trapping_instance = { + 'desc': 'Sample trapping instance', + 0: {'apside_1': [1, 2], 'apside_2': [3, 4]}, + 1: {'apside_1': [5, 6], 'apside_2': [7, 8]}, + 2: {'apside_1': [9, 10], 'apside_2': [11, 12]} + } + >>> reduced_instance = reduce_aps_dictionary(trapping_instance, 2) + >>> print(reduced_instance) + {'norb': 2, 'desc': 'Sample trapping instance', 0: {'apside_1': [1, 2], 'apside_2': [3, 4]}, 1: {'apside_1': [5, 6], 'apside_2': [7, 8]}} + """ + # Initialize the output dictionary TrappingInstanceOut = {} + + # Add the number of orbits to the output dictionary TrappingInstanceOut['norb'] = norb + + # Add the description to the output dictionary TrappingInstanceOut['desc'] = TrappingInstance['desc'] - for i in range(0,norb): + # Loop through the specified number of orbits and add them to the output dictionary + for i in range(norb): TrappingInstanceOut[i] = TrappingInstance[i] return TrappingInstanceOut - -def evaluate_clusters_polar_legacy(K,maxima=False,rank=False,perc=0.): - ''' - evaluate_clusters_polar - calculate statistics for clusters in polar coordinates - - inputs - ------------- - K : number of clusters - maxima : (boolean, False) if True, use the maximum value from the clusters - rank - perc - - - returns - ------------- - theta_n - clustermean - clusterstd_r - clusterstd_t - - - - ''' +def beane_criteria(K): + """Implement the criteria from Beane et al. (2024) + + works best for polar classifications""" # how many clusters? k = K.K - if (rank) & (perc==0.): - print('evaluate_clusters_polar: Perc must be >0.') - return np.nan,np.nan,np.nan,np.nan - + # Compute radii and theta values from clusters + rad_clusters = np.array([np.linalg.norm(K.clusters[i], axis=1) for i in range(k)]) + the_clusters = np.array([np.arctan2(np.abs(K.clusters[i][:, 1]), np.abs(K.clusters[i][:, 0])) for i in range(k)]) - # compute radii and theta values from clusters - rad_clusters = np.array([np.sum(np.array(K.clusters[i])*np.array(K.clusters[i]),axis=1)**0.5 for i in range(0,k)]) + # implement equation A1: the maximum angle from the bar for the clusters + thetadiff = np.max([np.arctan2(np.abs(K.mu[i][1]), np.abs(K.mu[i][0])) for i in range(k)]) - # for computing the theta values, can decide on a version with - # (legacy) or without (modern) folding + # if thetadiff < pi/8, consider the particle trapped - legacy = True + # implement equation A2: + clusterstd = np.sum([np.std(rad_clusters[i]) for i in range(k)]) + clustermean = np.sum([np.mean(rad_clusters[i]) for i in range(k)]) + tightness = clusterstd/clustermean - if legacy: - the_clusters = np.array([np.arctan(np.abs(np.array(K.clusters[i])[:,1])/np.abs(np.array(K.clusters[i])[:,0])) for i in range(0,k)]) - else: - the_clusters = np.array([np.arctan( (np.array(K.clusters[i])[:,1])/np.abs(np.array(K.clusters[i])[:,0])) for i in range(0,k)]) - - if maxima: - # use maxima + # if tightness is < 0.22, consider the particle trapped - clustermean = np.max([np.mean(rad_clusters[i]) for i in range(0,k)]) - - if legacy: - theta_n = np.max([np.abs(np.arctan( K.mu[i][1]/K.mu[i][0])) for i in range(0,k)]) - else: - theta_n = np.max([ np.arctan(np.abs(K.mu[i][1]/K.mu[i][0])) for i in range(0,k)]) - - if rank: - # use rank ordered - - organized_rad = np.array([rad_clusters[i][rad_clusters[i].argsort()] for i in range(0,k)]) - organized_the = np.array([the_clusters[i][the_clusters[i].argsort()] for i in range(0,k)]) - - clusterstd_r = np.max([np.percentile(organized_rad[i] - np.mean(rad_clusters[i]),perc) for i in range(0,k)]) - - # I think this needs an absolute value - clusterstd_t = np.max([np.percentile(organized_the[i] - np.mean(the_clusters[i]),perc) for i in range(0,k)]) - - else: - clusterstd_r = np.max([np.std(rad_clusters[i]) for i in range(0,k)]) - - if legacy: - clusterstd_t = np.max([np.std(the_clusters[i]) for i in range(0,k)]) - else: - clusterstd_t = np.max([np.abs(np.max(the_clusters[i])-np.min(the_clusters[i])) for i in range(0,k)]) - - - else: - - # not maxima - - if rank: - # use rank ordered - - organized_rad = np.array([rad_clusters[i][rad_clusters[i].argsort()] for i in range(0,k)]) - organized_the = np.array([the_clusters[i][the_clusters[i].argsort()] for i in range(0,k)]) - - clusterstd_r = np.mean([np.percentile(organized_rad[i] - np.mean(rad_clusters[i]),perc) for i in range(0,k)]) - clusterstd_t = np.mean([np.percentile(organized_the[i] - np.mean(the_clusters[i]),perc) for i in range(0,k)]) - - else: - clusterstd_r = np.mean([np.std(rad_clusters[i]) for i in range(0,k)]) - - if legacy: - clusterstd_t = np.mean([np.std(the_clusters[i]) for i in range(0,k)]) - else: - clusterstd_t = np.mean([np.abs(np.max(the_clusters[i])-np.min(the_clusters[i])) for i in range(0,k)]) - - - - clustermean = np.mean([np.mean(rad_clusters[i]) for i in range(0,k)]) - - # compute the mean of the cluster centers - if legacy: - theta_n = np.mean([np.abs(np.arctan( K.mu[i][1]/K.mu[i][0])) for i in range(0,k)]) - else: - theta_n = np.mean([ np.arctan(np.abs(K.mu[i][1]/K.mu[i][0])) for i in range(0,k)]) - - # return values - return theta_n,clustermean,clusterstd_r,clusterstd_t - - -def evaluate_clusters_polar(K,maxima=False,rank=False,perc=0.): - ''' - evaluate_clusters_polar - calculate statistics for clusters in polar coordinates - - inputs - ------------- - K - maxima - rank - perc - - - returns - ------------- - theta_n - clustermean - clusterstd_r - clusterstd_t + return thetadiff,tightness +def evaluate_clusters_polar(K, maxima=False, rank=False, perc=0.): + """ + Calculate statistics for clusters in polar coordinates. - ''' + This function evaluates the clustering results in polar coordinates (r, theta). + It computes various statistics such as mean, standard deviation, and optionally + ranks and percentiles. - # how many clusters? + Parameters + ---------- + K : KMeans + An instance of a K-means clustering result. + maxima : bool, optional + If True, calculate maximum quantities. If False, calculate average quantities. + Default is False. + rank : bool, optional + If True, use rank-ordered statistics. Default is False. + perc : float, optional + Percentage threshold for rank ordering. Default is 0. + + Returns + ------- + theta_n : float + Angle measure in the context of the clusters. + clustermean : float + The mean value of the clusters. + clusterstd_r : float + The standard deviation of the clusters in the radial direction. + clusterstd_t : float + The standard deviation of the clusters in the angular direction. + + Notes + ----- + This function computes the radii and theta values from the clusters, then calculates + either the maximum or average statistics based on the `maxima` parameter. It also + handles rank-ordered statistics if `rank` is True and `perc` is greater than 0. + + Examples + -------- + >>> K = kmeans.KMeans(k=2, X=ApsArray) + >>> K.find_centers() + >>> theta_n, clustermean, clusterstd_r, clusterstd_t = evaluate_clusters_polar(K) + """ + # Number of clusters k = K.K - if (rank) & (perc==0.): - print('evaluate_clusters_polar: Perc must be >0.') - return np.nan,np.nan,np.nan,np.nan + # Check if rank is True but perc is not set + if rank and perc == 0.: + raise SyntaxError('exptool.trapping.evaluate_clusters_polar: Perc must be >0.') - - # compute radii and theta values from clusters - rad_clusters = np.array([np.sum(np.array(K.clusters[i])*np.array(K.clusters[i]),axis=1)**0.5 for i in range(0,k)]) - the_clusters = np.array([np.arctan(np.abs(np.array(K.clusters[i])[:,1])/np.abs(np.array(K.clusters[i])[:,0])) for i in range(0,k)]) + # Compute radii and theta values from clusters + rad_clusters = np.array([np.linalg.norm(K.clusters[i], axis=1) for i in range(k)]) + the_clusters = np.array([np.arctan2(np.abs(K.clusters[i][:, 1]), np.abs(K.clusters[i][:, 0])) for i in range(k)]) if maxima: - # use maxima - - clustermean = np.max([np.mean(rad_clusters[i]) for i in range(0,k)]) - - #theta_n = np.max([abs(np.arctan(K.mu[i][1]/K.mu[i][0])) for i in range(0,k)]) - theta_n = np.max([np.arctan(np.abs(K.mu[i][1]/K.mu[i][0])) for i in range(0,k)]) + # Calculate maximum quantities + clustermean = np.max([np.mean(rad_clusters[i]) for i in range(k)]) + theta_n = np.max([np.arctan2(np.abs(K.mu[i][1]), np.abs(K.mu[i][0])) for i in range(k)]) if rank: - # use rank ordered - - organized_rad = np.array([rad_clusters[i][rad_clusters[i].argsort()] for i in range(0,k)]) - organized_the = np.array([the_clusters[i][the_clusters[i].argsort()] for i in range(0,k)]) - - clusterstd_r = np.max([np.percentile(organized_rad[i] - np.mean(rad_clusters[i]),perc) for i in range(0,k)]) - clusterstd_t = np.max([np.percentile(organized_the[i] - np.mean(the_clusters[i]),perc) for i in range(0,k)]) - + # Use rank-ordered statistics + organized_rad = np.array([np.sort(rad_clusters[i]) for i in range(k)]) + organized_the = np.array([np.sort(the_clusters[i]) for i in range(k)]) + clusterstd_r = np.max([np.percentile(organized_rad[i] - np.mean(rad_clusters[i]), perc) for i in range(k)]) + clusterstd_t = np.max([np.percentile(organized_the[i] - np.mean(the_clusters[i]), perc) for i in range(k)]) else: - clusterstd_r = np.max([np.std(rad_clusters[i]) for i in range(0,k)]) - clusterstd_t = np.max([np.std(the_clusters[i]) for i in range(0,k)]) - + clusterstd_r = np.max([np.std(rad_clusters[i]) for i in range(k)]) + clusterstd_t = np.max([np.std(the_clusters[i]) for i in range(k)]) else: - - # not maxima - + # Calculate average quantities if rank: - # use rank ordered - - organized_rad = np.array([rad_clusters[i][rad_clusters[i].argsort()] for i in range(0,k)]) - organized_the = np.array([the_clusters[i][the_clusters[i].argsort()] for i in range(0,k)]) - - clusterstd_r = np.mean([np.percentile(organized_rad[i] - np.mean(rad_clusters[i]),perc) for i in range(0,k)]) - clusterstd_t = np.mean([np.percentile(organized_the[i] - np.mean(the_clusters[i]),perc) for i in range(0,k)]) - + # Use rank-ordered statistics + organized_rad = np.array([np.sort(rad_clusters[i]) for i in range(k)]) + organized_the = np.array([np.sort(the_clusters[i]) for i in range(k)]) + clusterstd_r = np.mean([np.percentile(organized_rad[i] - np.mean(rad_clusters[i]), perc) for i in range(k)]) + clusterstd_t = np.mean([np.percentile(organized_the[i] - np.mean(the_clusters[i]), perc) for i in range(k)]) else: - clusterstd_r = np.mean([np.std(rad_clusters[i]) for i in range(0,k)]) - clusterstd_t = np.mean([np.std(the_clusters[i]) for i in range(0,k)]) - - clustermean = np.mean([np.mean(rad_clusters[i]) for i in range(0,k)]) - #theta_n = np.mean([abs(np.arctan(K.mu[i][1]/K.mu[i][0])) for i in range(0,k)]) - theta_n = np.mean([np.arctan(np.abs(K.mu[i][1]/K.mu[i][0])) for i in range(0,k)]) - - return theta_n,clustermean,clusterstd_r,clusterstd_t - + clusterstd_r = np.mean([np.std(rad_clusters[i]) for i in range(k)]) + clusterstd_t = np.mean([np.std(the_clusters[i]) for i in range(k)]) + clustermean = np.mean([np.mean(rad_clusters[i]) for i in range(k)]) + theta_n = np.mean([np.arctan2(np.abs(K.mu[i][1]), np.abs(K.mu[i][0])) for i in range(k)]) + return theta_n, clustermean, clusterstd_r, clusterstd_t -def process_kmeans_polar(ApsArray,indx=-1,k=2,maxima=False,rank=False,perc=0.): - ''' - # - # robust kmeans implementation - # - # -can be edited for speed - # -confined to two dimensions - # -computes trapping metrics in polar coordinates - - inputs - ---------- - ApsArray : the array of aps for an individual orbit - indx : a designation of the orbit, for use with multiprocessing - k : the number of clusters - maxima : calculate average (if False) or maximum (if True) quantities - mad : toggle median absolute deviation calculation +def process_kmeans_polar(ApsArray, indx=-1, k=2, maxima=False, rank=False, perc=0.): + """ + Perform robust K-means clustering on apsidal data in polar coordinates. + This function performs K-means clustering on the provided apsidal array, + computes trapping metrics in polar coordinates, and handles potential + edge cases where clusters may have very few points. - returns + Parameters ---------- - theta_n : (see explanation at beginning for definitions) - clustermean : - clusterstd_r : - clusterstd_theta : - kmeans_plus_flag : - - - - - ''' + ApsArray : array-like + The array of apsides for an individual orbit. Each element should contain + the (r, theta) coordinates of an apsis. + indx : int, optional + A designation of the orbit, for use with multiprocessing. Default is -1. + k : int, optional + The number of clusters to form. Default is 2. + maxima : bool, optional + Calculate average (if False) or maximum (if True) quantities. Default is False. + rank : bool, optional + Toggle ranking. Default is False. + perc : float, optional + Percentage threshold for ranking. Default is 0. + + Returns + ------- + theta_n : float + Some angle measure in the context of the clusters. + clustermean : float + The mean value of the clusters. + clusterstd_r : float + The standard deviation of the clusters in the radial direction. + clusterstd_theta : float + The standard deviation of the clusters in the angular direction. + kmeans_plus_flag : int + Indicator flag: 0 for successful basic K-means, 1 for successful K-means++, + 2 for failure in both K-means and K-means++. + + Notes + ----- + This implementation confines the clustering to two dimensions and includes + robustness checks to handle small cluster sizes. In case of failure in + basic K-means, it retries using the K-means++ initialization method. + + Examples + -------- + >>> ApsArray = np.array([[1.0, 0.0], [2.0, 1.0], [1.5, 0.5], [3.0, 1.5]]) + >>> theta_n, clustermean, clusterstd_r, clusterstd_theta, flag = process_kmeans_polar(ApsArray) + """ kmeans_plus_flag = 0 - K = kmeans.KMeans(k,X=ApsArray) + K = kmeans.KMeans(k, X=ApsArray) K.find_centers() - # add an evaluation for if a cluster ends up with only X members, here hard coded to 2 - - - # find the standard deviation of clusters - - # first, check to make sure no single-point clusters were detected - # set rejection threshold + # Minimum cluster size threshold min_cluster_size = 1 - try: + clustersize = np.array([np.array(K.clusters[c]).size / 2. for c in range(k)]) - clustersize = np.array([np.array(K.clusters[c]).size/2. for c in range(0,k)]) - - # eliminate + # Ensure no single-point clusters while np.min(clustersize) <= min_cluster_size: w = np.where(clustersize > min_cluster_size)[0] - new_aps = np.array([np.concatenate([np.array(K.clusters[x])[:,0] for x in w]),\ - np.concatenate([np.array(K.clusters[x])[:,1] for x in w])]).T + new_aps = np.array([np.concatenate([np.array(K.clusters[x])[:, 0] for x in w]), \ + np.concatenate([np.array(K.clusters[x])[:, 1] for x in w])]).T - K = kmeans.KMeans(k,X=new_aps) + K = kmeans.KMeans(k, X=new_aps) K.find_centers() - clustersize = np.array([np.array(K.clusters[c]).size/2. for c in range(0,k)]) - - theta_n,clustermean,clusterstd_r,clusterstd_t = \ - evaluate_clusters_polar(K,maxima=maxima,rank=rank,perc=perc) - + clustersize = np.array([np.array(K.clusters[c]).size / 2. for c in range(k)]) + theta_n, clustermean, clusterstd_r, clusterstd_theta = \ + evaluate_clusters_polar(K, maxima=maxima, rank=rank, perc=perc) - # failure on basic kmeans except: - K = kmeans.KPlusPlus(2,X=ApsArray) + # If basic K-means fails, try K-means++ + K = kmeans.KPlusPlus(k, X=ApsArray) K.init_centers() K.find_centers(method='++') kmeans_plus_flag = 1 try: - - clustersize = np.array([np.array(K.clusters[c]).size/2. for c in range(0,k)]) - + clustersize = np.array([np.array(K.clusters[c]).size / 2. for c in range(k)]) while np.min(clustersize) <= min_cluster_size: w = np.where(clustersize > min_cluster_size)[0] - new_aps = np.array([np.concatenate([np.array(K.clusters[x])[:,0] for x in w]),\ - np.concatenate([np.array(K.clusters[x])[:,1] for x in w])]).T + new_aps = np.array([np.concatenate([np.array(K.clusters[x])[:, 0] for x in w]), \ + np.concatenate([np.array(K.clusters[x])[:, 1] for x in w])]).T - K = kmeans.KPlusPlus(k,X=new_aps) + K = kmeans.KPlusPlus(k, X=new_aps) K.init_centers() K.find_centers(method='++') - clustersize = np.array([np.array(K.clusters[c]).size/2. for c in range(0,k)]) - - - theta_n,clustermean,clusterstd_r,clusterstd_t = \ - evaluate_clusters_polar(K,maxima=maxima,rank=rank,perc=perc) - + clustersize = np.array([np.array(K.clusters[c]).size / 2. for c in range(k)]) + theta_n, clustermean, clusterstd_r, clusterstd_theta = \ + evaluate_clusters_polar(K, maxima=maxima, rank=rank, perc=perc) - # failure mode for advanced kmeans - except: - - # - # would like a more intelligent way to diagnose - #if indx >= 0: - # print 'Orbit %i even failed in Kmeans++!!' %indx + except Exception: + # If both methods fail, set all outputs to NaN clusterstd_r = np.nan - clusterstd_t = np.nan + clusterstd_theta = np.nan clustermean = np.nan theta_n = np.nan kmeans_plus_flag = 2 - - - return theta_n,clustermean,clusterstd_r,clusterstd_t,kmeans_plus_flag - - + return theta_n, clustermean, clusterstd_r, clusterstd_theta, kmeans_plus_flag @@ -878,29 +840,43 @@ def process_kmeans(ApsArray,indx=-1,k=2,maxima=False,mad=False): +def transform_aps(ApsArray, BarInstance): + """ + Transform the apsides array into the bar frame of reference. + This function transforms the apsides array, aligning it with the bar frame as determined + by the BarInstance. The transformation is offloaded for clarity. -def transform_aps(ApsArray,BarInstance): - ''' - transform_aps : simple transformation for the aps array, offloaded for clarity. + Parameters + ---------- + ApsArray : np.ndarray + The array of apsides, where each row represents a time step and contains + [time, x_position, y_position]. + BarInstance : object + An instance that contains information about the bar's position and motion. + + Returns + ------- + np.ndarray + The transformed positions in the bar frame. The output array has the same number of rows + as ApsArray and two columns corresponding to the transformed x and y positions. + + Notes + ----- + This transformation assumes that the bar motion is in one direction. + """ - inputs - ------------------ - ApsArray : the array of apsides - BarInstance : + # Find the bar angle positions corresponding to the times in ApsArray + bar_positions = pattern.find_barangle(ApsArray[:, 0], BarInstance) - outputs - ------------------ - X : + # Initialize the output array for transformed positions + X = np.zeros([len(ApsArray[:, 1]), 2]) - stuck in one direction, watch out - ''' - bar_positions = pattern.find_barangle(ApsArray[:,0],BarInstance) - X = np.zeros([len(ApsArray[:,1]),2]) - X[:,0] = ApsArray[:,1]*np.cos(bar_positions) - ApsArray[:,2]*np.sin(bar_positions) - X[:,1] = -ApsArray[:,1]*np.sin(bar_positions) - ApsArray[:,2]*np.cos(bar_positions) - return X + # Apply the transformation to align with the bar frame + X[:, 0] = ApsArray[:, 1] * np.cos(bar_positions) - ApsArray[:, 2] * np.sin(bar_positions) + X[:, 1] = -ApsArray[:, 1] * np.sin(bar_positions) - ApsArray[:, 2] * np.cos(bar_positions) + return X def do_single_kmeans_step(TrappingInstanceDict,BarInstance,desired_time,\ @@ -1093,9 +1069,8 @@ def do_kmeans_dict(TrappingInstanceDict,BarInstance,\ norb = TrappingInstanceDict['norb'] nfamilies = len(criteria.keys()) if nfamilies == 0: - print('trapping.do_kmeans_dict: no families defined?') - #break - return + raise ValueError('exptool.trapping.do_kmeans_dict: no families defined?') + # set up final array trapping_array = np.zeros([nfamilies,norb,len(BarInstance.time)],dtype='i1') @@ -1463,7 +1438,7 @@ def do_kmeans_multi(TrappingInstanceDict,BarInstance,\ print('Total trapping calculation took {0:3.2f} seconds, or {1:3.2f} milliseconds per orbit.'.format(time.time()-t1, 1.e3*(time.time()-t1)/len(TrappingInstanceDict))) # go through the dictionary of trapping criteria and re-make the arrays - trapped = {} + trapped = dict() for nfam,family in enumerate(np.array(list(criteria.keys()))): @@ -1519,4 +1494,3 @@ def re_form_trapping_arrays(array,array_number): return net_array -#warnings.filterwarnings("ignore",category =RuntimeWarning) diff --git a/exptool/io/psp_to_hdf5.py b/exptool/io/psp_to_hdf5.py new file mode 100644 index 0000000..ed3c931 --- /dev/null +++ b/exptool/io/psp_to_hdf5.py @@ -0,0 +1,188 @@ +""" +draft conversion from PSP format to HDF5. + + +For each component group, there is a subgroup named 'header', which stores header information related to that component. This information may include various parameters and metadata. The header information is organized into nested groups and attributes within the 'header' subgroup. The structure of the header data may vary depending on the specific PSP file format. + +The main data associated with each component is stored in a dataset named 'phasespace' within the component group. This dataset is an Nx8 array, where N represents the number of particles. Each row of the dataset corresponds to a particle and contains the following information in this order: mass, x-coordinate, y-coordinate, z-coordinate, x-velocity, y-velocity, z-velocity, and potential energy. + + +Example usage: +import argparse + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + prog='PSP2HDF5', + description='Convert OUTPSN files to HDF5 format.')#epilog='Text at the bottom of help') + + parser.add_argument('filename',help='the file to be converted') + + args = parser.parse_args() + + HDFConverter(args.filename) + + + +# example usage +# to run, do PSP2HDF5 OUT.run0.00000 + +import h5py +outputfilename = 'OUT.run0.00000.h5' + +# how is the global header information saved? +# only time is saved +f = h5py.File(outputfilename, 'r') + +# if the component is called 'halo', then +f['halo/phasespace'] # is the dataset, (Nx8), with (mass,x,y,z,vx,vy,vz,potential) for each particle + +# the header data is saved in a group, +f['halo/header'] +# and then each group may have subgroups or attributes. typically, you will have +f['halo/header/parameters'] +print(f['halo/header/parameters'].attrs.keys()) +print(f['halo/header/parameters'].attrs['nEJwant']) + +f['halo/header/force'] +print(f['halo/header/force'].attrs['id']) +print(f['halo/header/force/parameters'].attrs.keys()) +for key in f['halo/header/force/parameters'].attrs.keys(): + print(key,f['halo/header/force/parameters'].attrs[key]) + + +""" + + + +import numpy as np +import h5py +from . import particle + +class HDFConverter(): + """ + HDFConverter class for converting custom PSP (Particle Simulation Program) files to HDF5 format. + + This class allows you to convert data from PSP format into HDF5 format, which is a versatile and efficient data storage format. + + Parameters: + filename (str): The name of the PSP input file to be converted. + comp (str): Optional. The specific component to convert. If provided, only the data for the specified component will be converted. + verbose (int): Optional. Verbosity level for printing progress and messages during conversion. + + Attributes: + filename (str): The name of the PSP input file. + """ + + def __init__(self, filename, comp=None, verbose=0): + """ + Initialize the HDFConverter instance. + + Args: + filename (str): The name of the PSP input file to be converted. + comp (str, optional): The specific component to convert. If provided, only the data for the specified component will be converted. + verbose (int, optional): Verbosity level for printing progress and messages during conversion. + """ + self.filename = filename + self.comp = comp + self.verbose = verbose + + # Start the conversion process + self.convert_psp_to_hdf5(self.filename, comp=self.comp, verbose=self.verbose) + + def convert_psp_to_hdf5(self): + """ + Convert a PSP input file to HDF5 format. + + Uses the filename stored in self.filename during initialization. + + """ + # Define the output file name + outputfilename = self.filename + '.h5' + + if verbose > 0: + print(f"Converting {inputfilename} to {outputfilename}") + + # Open the PSP input file and extract components + O = particle.Input(self.filename) + comps = list(O.header.keys()) + + + # Create a new HDF5 file for storing the converted data + f = h5py.File(outputfilename, 'w') + + # Store the simulation time as an attribute + f['time'] = O.time + + for comp in comps: + if verbose > 0: + print(f"Processing component: {comp}") + + # Create a group for each component + f.create_group(comp) + + # Print the header information for the component + self.print_component_header(f, O, comp) + + # Create and store the phase space data for the component + self.make_phasespace(f, comp) + + # Close the HDF5 file + f.close() + + if verbose > 0: + print(f"Conversion complete: {outputfilename}") + + def print_component_header(self, f, O, comp): + """ + Print header information for a component to an HDF5 file. + + Args: + f (h5py.Group): The HDF5 group to store the header information. + O (particle.Input): The PSP input object. + comp (str): The name of the component. + """ + f[comp].create_group('header') + for key in O.header[comp].keys(): + # Check for nested dictionary levels + try: + for subkey in O.header[comp][key].keys(): + # Check for further nested levels + try: + for subsubkey in O.header[comp][key][subkey].keys(): + # Create attributes for the deepest level + try: + f['{}/header/{}/{}'.format(comp, key, subkey)].attrs.create(subsubkey, O.header[comp][key][subkey][subsubkey]) + except (KeyError, AttributeError, TypeError): + # Create subgroups if necessary + f['{}/header/{}/{}'.format(comp, key, subkey)].create_group(subsubkey) + f['{}/header/{}/{}'.format(comp, key, subkey)].attrs.create(subsubkey, O.header[comp][key][subkey][subsubkey]) + except (KeyError, AttributeError, TypeError): + # Create attributes for the intermediate level + try: + f['{}/header/{}/{}'.format(comp, key, subkey)].attrs.create(subkey, O.header[comp][key][subkey]) + except (KeyError, AttributeError, TypeError): + # Create subgroups if necessary + f['{}/header'.format(comp)].create_group(key) + f['{}/header/{}'.format(comp, key)].attrs.create(subkey, O.header[comp][key][subkey]) + except (KeyError, AttributeError, TypeError): + # Create attributes for the top-level header + f['{}/header'.format(comp)].attrs.create(key, O.header[comp][key]) + + def make_phasespace(self, f, comp): + """ + Convert and store phase space data for a component in an HDF5 file. + + Args: + f (h5py.Group): The HDF5 group to store the phase space data. + comp (str): The name of the component. + + Uses the filename stored in self.filename during initialization. + """ + # Read data from the PSP input file + O1 = particle.Input(self.filename, comp) + + # Create a phase space array + PS = np.array([O1.data['m'], O1.data['x'], O1.data['y'], O1.data['z'], O1.data['vx'], O1.data['vy'], O1.data['vz'], O1.data['potE']]).T + + # Store the phase space data as a dataset + f[comp].create_dataset('phasespace', data=PS) diff --git a/exptool/utils/kmeans.py b/exptool/utils/kmeans.py index a4c79b9..f843bbf 100644 --- a/exptool/utils/kmeans.py +++ b/exptool/utils/kmeans.py @@ -1,119 +1,97 @@ -# -# kmeans.py -# -# robust implementation of kmeans based on -# https://datasciencelab.wordpress.com +""" +kmeans.py -import random -import numpy as np +purpose-built, robust implementation of kmeans based on +https://datasciencelab.wordpress.com -import matplotlib.pyplot as plt -import matplotlib +MSP 19 May 2024 Improve documentation; code cleanup +""" -class KMeans(): - ''' - class to implement K-means in Python +import numpy as np +import random +class KMeans: + """ + A class to implement K-means clustering in Python. + """ - ''' - def __init__(self, K, X=None, N=0): - ''' - initialize kmeans - - - inputs - -------------- - self : KMeans class - K : number of clusters - X : array of observations - N : error guard - - - returns - ------------- - self : Kmeans class - - - ''' - + """ + Initialize KMeans. + + Parameters + ---------- + K : int + Number of clusters. + X : array-like, optional + Array of observations. + N : int, optional + Number of points (needed if X is not provided). + + Raises + ------ + Exception + If no data is provided and N is not specified. + """ self.K = K - try: - - tmp = len(X) + if X is not None: self.X = X self.N = len(X) - - except: - + else: if N == 0: - raise Exception("kmeans.KMeans: If no data is provided, \ - a parameter N (number of points) is needed") + raise Exception("kmeans.KMeans: If no data is provided, a parameter N (number of points) is needed") else: self.N = N self.X = self._init_board_gauss(N, K) self.mu = None + self.oldmu = None self.clusters = None self.method = None - def _init_board_gauss(self, N, k): - ''' - _init_board_gauss - initialize the guess points - - - inputs - ------------------- - self - N - k - - - returns - ------------------ - self - X : randomly partitioned clusters - - - ''' - - # number of points to put in each cluster - n = float(N)/k - + def _init_board_gauss(self, N, K): + """ + Initialize the guess points using a Gaussian distribution. + + Parameters + ---------- + N : int + Number of points. + K : int + Number of clusters. + + Returns + ------- + np.array + Randomly partitioned clusters. + """ + n = float(N) / K X = [] - # set up - for i in range(k): - - c = (random.uniform(-1,1), random.uniform(-1,1)) - - s = random.uniform(0.05,0.15) - - # just reflecting--but this means that the clusters are forced to ahve the same number of points - x = [] - while len(x) < n: - - a,b = np.array([np.random.normal(c[0],s),np.random.normal(c[1],s)]) - - # Continue drawing points from the distribution in the range [-1,1] - if abs(a) and abs(b)<1: - x.append([a,b]) - - X.extend(x) + for i in range(K): + c = (random.uniform(-1, 1), random.uniform(-1, 1)) + s = random.uniform(0.05, 0.15) + + cluster_points = [] + while len(cluster_points) < n: + a, b = np.array([np.random.normal(c[0], s), np.random.normal(c[1], s)]) + if abs(a) < 1 and abs(b) < 1: + cluster_points.append([a, b]) + X.extend(cluster_points) X = np.array(X)[:N] return X - def _cluster_points(self): + """ + Assign each point to the nearest cluster center. + """ mu = self.mu - clusters = {} + clusters = {} for x in self.X: - bestmukey = min([(i[0], np.linalg.norm(x-mu[i[0]])) \ - for i in enumerate(mu)], key=lambda t:t[1])[0] + bestmukey = min([(i[0], np.linalg.norm(x - mu[i[0]])) for i in enumerate(mu)], key=lambda t: t[1])[0] try: clusters[bestmukey].append(x) except KeyError: @@ -121,99 +99,140 @@ def _cluster_points(self): self.clusters = clusters def _reevaluate_centers(self): - ''' - _reevaluate_centers - draw new centers based on which center they are closest to - NOTE that this can create asymmetric cluster sizes - - ''' + """ + Compute new cluster centers based on the current cluster assignments. + """ clusters = self.clusters newmu = [] - keys = sorted(self.clusters.keys()) + keys = sorted(clusters.keys()) for k in keys: - newmu.append(np.mean(clusters[k], axis = 0)) + newmu.append(np.mean(clusters[k], axis=0)) self.mu = newmu def _has_converged(self): - ''' - _has_converged - check to see whether clusters change from one step to the next. if not, declare convergence! - - - ''' + """ + Check if the algorithm has converged. + + Returns + ------- + bool + True if the cluster centers do not change, False otherwise. + """ K = len(self.oldmu) - return(set([tuple(a) for a in self.mu]) == \ - set([tuple(a) for a in self.oldmu])\ - and len(set([tuple(a) for a in self.mu])) == K) + return (set([tuple(a) for a in self.mu]) == set([tuple(a) for a in self.oldmu]) and + len(set([tuple(a) for a in self.mu])) == K) - def find_centers(self, method='random',nitermax=1000): - ''' - find_centers - iteratively select new centers - - inputs - ---------------- - - - returns - --------------- - - - ''' - + def find_centers(self, method='random', nitermax=1000): + """ + Find the cluster centers using the specified method. + + Parameters + ---------- + method : str, optional + Method to initialize the cluster centers ('random' or '++', default is 'random'). + nitermax : int, optional + Maximum number of iterations (default is 1000). + + Returns + ------- + None + + Raises + ------ + ValueError + If the specified method is not supported. + """ self.method = method X = self.X K = self.K - - #self.oldmu = random.sample(X, K) - # draw K samples from the array of clusters - self.oldmu = random.sample(list(X), K) - - - if method != '++': + if method == '++': + # Initialize using K-means++ + self.mu = self._init_kmeans_plusplus() + elif method == 'random': # Initialize to K random centers - - # this has a python2/3 compatibility issue - #self.mu = random.sample(X, K) self.mu = random.sample(list(X), K) + else: + raise ValueError("Unsupported method: {}".format(method)) - # put in a guard against + # Initialize oldmu to ensure first iteration runs + # Use values guaranteed to be different from initial mu + self.oldmu = [np.array([float('inf')] * len(X[0])) for _ in range(K)] + iter = 0 - - while (not self._has_converged()) & (iter= r)[0][0] - return(self.X[ind]) + return self.X[ind] def init_centers(self): - #self.mu = random.sample(self.X, 1) + """ + Initialize the cluster centers using K-means++ initialization. + + Returns + ------- + None + """ self.mu = random.sample(list(self.X), 1) while len(self.mu) < self.K: self._dist_from_centers() self.mu.append(self._choose_next_center()) -