Fix behavior at edges in field_derivs, potential_derivs

leon-vv · Nov 10, 2024 · 53835a4 · 53835a4
1 parent 5e81991
commit 53835a4
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/traceon/backend/radial.c b/traceon/backend/radial.c
@@ -315,12 +315,17 @@ potential_radial_derivs(double point[2], double *z_inter, double *coeff_p, size_
 	double r = point[0], z = point[1];
 	double z0 = z_inter[0], zlast = z_inter[N_z-1];
 
-	if(!(z0 < z && z < zlast)) {
+	if(!(z0 <= z && z <= zlast)) {
 		return 0.0;
 	}
 
 	double dz = z_inter[1] - z_inter[0];
 	int index = (int) ( (z-z0)/dz );
+
+	// Guard against roundoff issues 0 <= index < N_z
+	index = MAX(0, index);
+	index = MIN(index, N_z-1);
+
 	double diffz = z - z_inter[index];
 
 	double (*C)[6] = &coeff[index][0];
@@ -342,13 +347,17 @@ field_radial_derivs(double point[3], double field[3], double *z_inter, double *c
 	double r = norm_2d(point[0], point[1]), z = point[2];
 	double z0 = z_inter[0], zlast = z_inter[N_z-1];
 
-	if(!(z0 < z && z < zlast)) {
+	if(!(z0 <= z && z <= zlast)) {
 		field[0] = 0.0; field[1] = 0.0; field[2] = 0.0;
 		return;
 	}
 
 	double dz = z_inter[1] - z_inter[0];
 	int index = (int) ( (z-z0)/dz );
+	// Guard against roundoff issues 0 <= index < N_z
+	index = MAX(0, index);
+	index = MIN(index, N_z-1);
+
 	double diffz = z - z_inter[index];
 
 	double (*C)[6] = &coeff[index][0];

diff --git a/traceon/backend/three_d.c b/traceon/backend/three_d.c
@@ -204,10 +204,14 @@ potential_3d_derivs(double point[3], double *zs, double *coeffs_p, size_t N_z) {
 
 	double xp = point[0], yp = point[1], zp = point[2];
 
-	if (!(zs[0] < zp && zp < zs[N_z-1])) return 0.0;
-
+	if (!(zs[0] <= zp && zp <= zs[N_z-1])) return 0.0;
+	
 	double dz = zs[1] - zs[0];
 	int index = (int) ((zp-zs[0])/dz);
+	// Guard against roundoff issues 0 <= index < N_z
+	index = MAX(0, index);
+	index = MIN(index, N_z-1);
+
 
 	double z_ = zp - zs[index];
 
@@ -241,11 +245,14 @@ field_3d_derivs(double point[3], double field[3], double *restrict zs, double *r
 
 	field[0] = 0.0, field[1] = 0.0, field[2] = 0.0;
 
-	if (!(zs[0] < zp && zp < zs[N_z-1])) return;
+	if (!(zs[0] <= zp && zp <= zs[N_z-1])) return;
 
 	double dz = zs[1] - zs[0];
 	int index = (int) ((zp-zs[0])/dz);
-
+	// Guard against roundoff issues 0 <= index < N_z
+	index = MAX(0, index);
+	index = MIN(index, N_z-1);
+
 	double z_ = zp - zs[index];
 
 	double A[NU_MAX][M_MAX], B[NU_MAX][M_MAX];