Skip to content

Commit

Permalink
did a quick parallelization for the c++ functions. Still seems slower…
Browse files Browse the repository at this point in the history 
… than it should be so might need additional debugging.
  • Loading branch information
@akaptano
akaptano committed Aug 23, 2024
1 parent e833224 commit d319db4
Show file tree
Hide file tree
Showing 4 changed files with 52 additions and 33 deletions.
6 changes: 3 additions & 3 deletions examples/2_Intermediate/tokamak_exact.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,14 @@
nphi = 4 # nphi = ntheta >= 64 needed for accurate full-resolution runs
ntheta = nphi
else:
nphi = 64 # nphi = ntheta >= 64 needed for accurate full-resolution runs
nphi = 16 # nphi = ntheta >= 64 needed for accurate full-resolution runs
ntheta = nphi
Nx = 50 # cartesian bricks but note that we are not modelling the cubic geometry!
Nx = 40 # cartesian bricks but note that we are not modelling the cubic geometry!
Ny = Nx
Nz = Nx

coff = 0.1 # PM grid starts offset ~ 10 cm from the plasma surface
poff = 0.5 # PM grid end offset ~ 15 cm from the plasma surface
poff = 0.25 # PM grid end offset ~ 15 cm from the plasma surface
input_name = 'input.circular_tokamak'

# Read in the plas/ma equilibrium file
Expand Down
6 changes: 5 additions & 1 deletion src/simsopt/geo/permanent_magnet_grid.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1034,12 +1034,16 @@ def _optimization_setup(self):
self.b_obj = - self.Bn.reshape(self.nphi * self.ntheta)

# Compute geometric factor with the C++ routine
self.A_obj = cub.Acube(
import time
t1 = time.time()
self.A_obj = sopp.Acube(
np.ascontiguousarray(self.plasma_boundary.gamma().reshape(-1, 3)),
np.ascontiguousarray(self.pm_grid_xyz),
np.ascontiguousarray(self.plasma_boundary.unitnormal().reshape(-1, 3)),
self.dims, self.phiThetas
)
t2 = time.time()
print('Acube took ', t2 - t1, ' seconds')

print('# points = ',len(np.ascontiguousarray(self.plasma_boundary.gamma().reshape(-1, 3))))
print('# mag points = ', len(np.ascontiguousarray(self.pm_grid_xyz)))
Expand Down
58 changes: 36 additions & 22 deletions src/simsoptpp/Bcube_nonVec.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ Array Pd(double phi, double theta) {
return P;
}

Array iterate_over_corners(Array corner, Array x, Array y, Array z) {
Array iterate_over_corners(Array& corner, Array& x, Array& y, Array& z) {
int i = corner[0], j = corner[1], k = corner[2];
double summa = std::pow(-1, i+j+k);
double rijk = std::sqrt(x[i]*x[i] + y[j]*y[j] + z[k]*z[k]);
Expand All @@ -57,7 +57,7 @@ Array iterate_over_corners(Array corner, Array x, Array y, Array z) {
return summa * h;
}

Array Hd_i_prime(Array r, Array dims) {
Array Hd_i_prime(Array& r, Array& dims) {

Array H = xt::zeros<double>({3, 3});

Expand Down Expand Up @@ -89,18 +89,22 @@ Array Hd_i_prime(Array r, Array dims) {
}


Array B_direct(Array points, Array magPos, Array M, Array dims, Array phiThetas) {
Array B_direct(Array& points, Array& magPos, Array& M, Array& dims, Array& phiThetas) {
int N = points.shape(0);
int D = M.shape(0);

Array B = xt::zeros<double>({N, 3});
#pragma omp parallel for
for (int n = 0; n < N; ++n) {
double x = points(n, 0);
double y = points(n, 1);
double z = points(n, 2);
for (int d = 0; d < D; ++d) {
Array P = Pd(phiThetas(d,0), phiThetas(d,1));
Array r = xt::zeros<double>({3});
r[0] = points(n,0) - magPos(d,0);
r[1] = points(n,1) - magPos(d,1);
r[2] = points(n,2) - magPos(d,2);
r[0] = x - magPos(d,0);
r[1] = y - magPos(d,1);
r[2] = z - magPos(d,2);
Array r_loc = xt::zeros<double>({3});
Array M_loc = xt::zeros<double>({3});
for (int i = 0; i < 3; ++i) {
Expand All @@ -112,33 +116,34 @@ Array B_direct(Array points, Array magPos, Array M, Array dims, Array phiThetas)

Array H = Hd_i_prime(r_loc, dims);

double tx = heaviside(dims[0]/2 - std::abs(r_loc[0]),0.5);
double ty = heaviside(dims[1]/2 - std::abs(r_loc[1]),0.5);
double tz = heaviside(dims[2]/2 - std::abs(r_loc[2]),0.5);
double tx = heaviside(dims[0]/2 - std::abs(r_loc[0]), 0.5);
double ty = heaviside(dims[1]/2 - std::abs(r_loc[1]), 0.5);
double tz = heaviside(dims[2]/2 - std::abs(r_loc[2]), 0.5);
double tm = 2*tx*ty*tz;

Array B_loc = xt::zeros<double>({3});
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < 3; ++j) {
B_loc(n,i) += mu0 * (H(i,j) * M_loc[j] + tm * M_loc[i]);
B_loc(n, i) += mu0 * (H(i,j) * M_loc[j] + tm * M_loc[i]);
}
}
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < 3; ++j) {
B(n,i) += P(j,i) * B_loc[j];
B(n, i) += P(j,i) * B_loc[j];
}
}
}
}
return B;
}

Array Bn_direct(Array points, Array magPos, Array M, Array norms, Array dims, Array phiThetas) {
Array Bn_direct(Array& points, Array& magPos, Array& M, Array& norms, Array& dims, Array& phiThetas) {
int N = points.shape(0);

Array B = B_direct(points, magPos, M, dims, phiThetas);
Array Bn = xt::zeros<double>({N});

#pragma omp parallel for
for (int n = 0; n < N; ++n) {
for (int i = 0; i < 3; ++ i) {
Bn[n] += B(n,i) * norms(n,i);
Expand All @@ -148,7 +153,7 @@ Array Bn_direct(Array points, Array magPos, Array M, Array norms, Array dims, Ar
}


Array gd_i(Array r_loc, Array n_i_loc, Array dims) {
Array gd_i(Array& r_loc, Array& n_i_loc, Array& dims) {
double tx = heaviside(dims[0]/2 - std::abs(r_loc[0]),0.5);
double ty = heaviside(dims[1]/2 - std::abs(r_loc[1]),0.5);
double tz = heaviside(dims[2]/2 - std::abs(r_loc[2]),0.5);
Expand All @@ -170,19 +175,23 @@ Array gd_i(Array r_loc, Array n_i_loc, Array dims) {
return mu0 * g_loc;
}

Array Acube(Array points, Array magPos, Array norms, Array dims, Array phiThetas) {
Array Acube(Array& points, Array& magPos, Array& norms, Array& dims, Array& phiThetas) {
int N = points.shape(0);
int D = magPos.shape(0);
double magVol = dims[0] * dims[1] * dims[2];

Array A = xt::zeros<double>({N,3*D});
Array A = xt::zeros<double>({N, 3*D});
#pragma omp parallel for
for (int n = 0; n < N; ++n) {
double x = points(n, 0);
double y = points(n, 1);
double z = points(n, 2);
for (int d = 0; d < D; ++d) {
Array P = Pd(phiThetas(d,0), phiThetas(d,1));
Array r = xt::zeros<double>({3});
r[0] = points(n,0) - magPos(d,0);
r[1] = points(n,1) - magPos(d,1);
r[2] = points(n,2) - magPos(d,2);
// double rx = x - magPos(d,0);
// double ry = y - magPos(d,1);
// double rz = z - magPos(d,2);
Array r_loc = xt::zeros<double>({3});
Array n_loc = xt::zeros<double>({3});
for (int i = 0; i < 3; ++i) {
Expand All @@ -191,6 +200,9 @@ Array Acube(Array points, Array magPos, Array norms, Array dims, Array phiThetas
n_loc[i] += P(i,j) * norms(n,j);
}
}
// double rlocx = P(0, 0) * rx + P(0, 1) * ry + P(0, 2) * rz;
// double rlocy = P(1, 0) * rx + P(1, 1) * ry + P(1, 2) * rz;
// double rlocz = P(2, 0) * rx + P(2, 1) * ry + P(2, 2) * rz;
Array g_loc = gd_i(r_loc, n_loc, dims);
Array g = xt::zeros<double>({3});
// double magVol = dims(d,0) * dims(d,1) * dims(d,2);
Expand All @@ -199,25 +211,27 @@ Array Acube(Array points, Array magPos, Array norms, Array dims, Array phiThetas
g[i] += (P(j,i) * g_loc[j]) / magVol;
}
}
A(n,3*d) = g[0];
A(n,3*d + 1) = g[1];
A(n,3*d + 2) = g[2];
A(n, 3*d) = g[0];
A(n, 3*d + 1) = g[1];
A(n, 3*d + 2) = g[2];
}
}
return A;
}

Array Bn_fromMat(Array points, Array magPos, Array M, Array norms, Array dims, Array phiThetas) {
Array Bn_fromMat(Array& points, Array& magPos, Array& M, Array& norms, Array& dims, Array& phiThetas) {
int N = points.shape(0);
int D = M.shape(0);
Array A = Acube(points, magPos, norms, dims, phiThetas);
Array Ms = xt::zeros<double>({3*D});
#pragma omp parallel for
for (int d = 0; d < D; ++d) {
for (int i = 0; i < 3; ++i) {
Ms(3*d + i) = M(d,i);
}
}
Array Bn = xt::zeros<double>({N});
#pragma omp parallel for
for (int n = 0; n < N; ++n) {
for (int d = 0; d < 3*D; ++d) {
Bn[n] += A(n,d) * Ms[d];
Expand Down
15 changes: 8 additions & 7 deletions src/simsoptpp/Bcube_nonVec.h
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
#pragma once

#include <cmath>
#include <iostream>
Expand All @@ -8,16 +9,16 @@ double heaviside(double x1, double x2);

Array Pd(double phi, double theta);

Array iterate_over_corners(Array corner, Array x, Array y, Array z);
Array iterate_over_corners(Array& corner, Array& x, Array& y, Array& z);

Array Hd_i_prime(Array r, Array dims);
Array Hd_i_prime(Array& r, Array& dims);

Array B_direct(Array points, Array magPos, Array M, Array dims, Array phiThetas);
Array B_direct(Array& points, Array& magPos, Array& M, Array& dims, Array& phiThetas);

Array Bn_direct(Array point, Array magPos, Array M, Array norms, Array dims, Array phiThetas);
Array Bn_direct(Array& point, Array& magPos, Array& M, Array& norms, Array& dims, Array& phiThetas);

Array gd_i(Array r_loc, Array n_i_loc, Array dims);
Array gd_i(Array& r_loc, Array& n_i_loc, Array& dims);

Array Acube(Array points, Array magPos, Array norms, Array dims, Array phithetas);
Array Acube(Array& points, Array& magPos, Array& norms, Array& dims, Array& phithetas);

Array Bn_fromMat(Array points, Array magPos, Array M, Array norms, Array dims, Array phiThetas);
Array Bn_fromMat(Array& points, Array& magPos, Array& M, Array& norms, Array& dims, Array& phiThetas);

0 comments on commit d319db4

Please sign in to comment.