MagnetOptimization/script.py at main · HeartLab-McMaster/MagnetOptimization · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
# %% [markdown]
# ## Magnet optimization

# %%
from scipy.spatial.transform import Rotation as R
import numpy as np
import cvxpy as cp
import matplotlib.pyplot as plt
import magpylib as magpy


# %%
Br = 1.43 # (T) Residual flux density for N42
mu_0 = 4 * np.pi * 10**-7 # (H/m) Permeability of free space
l = 5.08e-2 # (m) Length of cube magnet
Volume = l ** 3 # (m^3)
moment = Br * Volume / mu_0 # (A m^2)
j = Br / mu_0 # (A/m)

# %%
Volume = 3.0e-3 ** 3
moment_target = Br * Volume / mu_0

# %%
target = np.array([0.15, 0.15, 0.45]) # target position is at 40 cm above the origin
workspace_length = 0.3 # workspace is a cube of 20 cm side length
mt = np.array([moment_target, 0, 0])

# %%
# return the magnetic field generated by a magnet at position p and orientation r
def generate_random_pose() -> tuple[np.ndarray, np.ndarray]:
    # generate a random pose
    r = R.random()
    p = np.random.rand(3) * workspace_length
    return p, r.as_matrix()

# %%
def B(p_i: np.ndarray, dm_i: np.ndarray):
  r_i = target - p_i
  r_i_hat = r_i / np.linalg.norm(r_i)
  return mu_0 * moment / (4 * np.pi * np.linalg.norm(r_i) ** 3) * ((3 * np.outer(r_i_hat, r_i_hat) - np.eye(3)) @ dm_i)

def F(p_i: np.ndarray, dm_i: np.ndarray):
  r_i = target - p_i
  r_i_hat = r_i / np.linalg.norm(r_i)
  return 3 * mu_0 * moment / (4 * np.pi * np.linalg.norm(r_i) ** 4) \
    * np.dot(
      np.outer(dm_i, r_i_hat) +
      np.outer(r_i_hat, dm_i) -
      ((5 * np.outer(r_i_hat, r_i_hat) - np.eye(3)) * np.dot(dm_i, r_i_hat))
      , mt)

def Jb(p_i: np.ndarray, dm_i: np.ndarray):
  r_i = target - p_i
  r_i_hat = r_i / np.linalg.norm(r_i)
  return mu_0 * moment / (4 * np.pi * np.linalg.norm(r_i) ** 3) * ((3 * np.outer(r_i_hat, r_i_hat) - np.eye(3)) @ dm_i)

def Jf(p_i: np.ndarray, dm_i: np.ndarray):
  r_i = target - p_i
  r_i_hat = r_i / np.linalg.norm(r_i)
  return 3 * mu_0 * moment / (4 * np.pi * np.linalg.norm(r_i) ** 4) \
    * np.dot(
      np.outer(dm_i, r_i_hat) +
      np.outer(r_i_hat, dm_i) -
      ((5 * np.outer(r_i_hat, r_i_hat) - np.eye(3)) * np.dot(dm_i, r_i_hat))
      , mt)

# %%
m = 500 # Number of random poses
K = 8 # Selection budget
d = 2 # Number of divisions for angles
n = d ** K

# %%
# Generating all combinations of angles
lins  = [np.linspace(0, 1.5*np.pi, d) for i in range(K)]
# lins.append(np.linspace(0, 2*np.pi, d) + np.pi/4)
angles = np.array(np.meshgrid(*lins)).T.reshape(-1, K)

# %%
# S is an array of tuples, each tuple contains a position and a rotation matrix
S = [generate_random_pose() for i in range(m)]

# %%
Bs = []
Fs = []
for p, r in S:
  m_i = np.array([0, 0, moment]) # all magnets having north pole facing upwards
  Bs.append(np.linalg.norm(B(p, m_i)))
  Fs.append(np.linalg.norm(F(p, m_i)))

Bmax = np.partition(Bs, -K)[-K:].sum() # Sum of the norms of K highest fields
Fmax = np.partition(Fs, -K)[-K:].sum() # Sum of the norms of K highest forces

# %%
# Initizaling A
A = np.zeros((n, K, m, 6, 6))

for t, theta in enumerate(angles):
  for i in range(K):
    for j, (p, r) in enumerate(S):
      dmagnetization = r.dot([- np.sin(theta[i]), np.cos(theta[i]), 0])
      J = np.concatenate([Jb(p, dmagnetization)/Bmax, Jf(p, dmagnetization)/Fmax])
      A[t, i, j, :, :] = np.outer(J, J)

# %%
def A_operator(X, t):
  return cp.sum([X[i][j] * A[t, i, j] for i in range(K) for j in range(m)])

# %%
X = cp.Variable(shape=(K, m))
t = cp.Variable(1)

alpha = 0
# obj = cp.Maximize(t + alpha * cp.atoms.lambda_min(f_operator(X)))
obj = cp.Maximize(t)
cons1 = X >= 0.0
cons2 = X <= 1.0
cons4 = cp.sum(X) == K # sum of all elements is K
cons5 = cp.sum(X, axis=1) == 1.0 # sum of each row is 1
cons6 = cp.sum(X, axis=0) <= 1.0 # sum of each col is le 1
cons7 = t >= 0.0
constraints = [cons1, cons2, cons5, cons6]
for i in range(n):
  constraints.append(t <= cp.atoms.lambda_min(A_operator(X, i)))
prob = cp.Problem(obj, constraints)

# %%
# prob.solve(verbose=True)

# %%
tol = 1.0e-5
prob.solve(verbose=True, solver=cp.CLARABEL, tol_gap_abs=tol, tol_gap_rel=tol, tol_feas=tol)

# %%
print("Status: ", prob.status)
# print("Solution x = ", X.value)
print("Solution t = ", t.value)


# %%
import pickle
dp = dict(S=S, X=X, t=t, params={"m": m, "K": K, "d": d, "n": n})

with open("runs/checkpoint_d5.pkl", "wb") as cp_file:
    pickle.dump(dp, cp_file)