Skip to content

Commit 38711f4

Browse files
A-HayasakaA-Hayasaka
authored
Merge pull request #21 from istellartech/fix_jacobian_format
Fix jacobian format
2 parents e0b0a84 + bbf7e24 commit 38711f4

File tree

1 file changed

+156
-91
lines changed

1 file changed

+156
-91
lines changed

‎constraints.py

Lines changed: 156 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
import sys
2727
import numpy as np
2828
from numpy.linalg import norm
29-
from scipy import sparse
3029
from numba import jit
3130
from utils import *
3231
from USStandardAtmosphere import *
@@ -115,36 +114,60 @@ def equality_jac_init(xdict, pdict, unitdict, condition):
115114
jac = {}
116115

117116
if condition["OptimizationMode"] == "Payload":
118-
rowcol_pos = (range(0, 3), range(0, 3))
119-
rowcol_vel = (range(3, 6), range(0, 3))
120-
rowcol_quat = (range(6, 10), range(0, 4))
121-
122-
jac["position"] = sparse.coo_matrix(
123-
([1.0] * 3, rowcol_pos), shape=(10, pdict["M"] * 3)
124-
)
125-
jac["velocity"] = sparse.coo_matrix(
126-
([1.0] * 3, rowcol_vel), shape=(10, pdict["M"] * 3)
127-
)
128-
jac["quaternion"] = sparse.coo_matrix(
129-
([1.0] * 4, rowcol_quat), shape=(10, pdict["M"] * 4)
130-
)
117+
jac["position"] = {
118+
"coo": [
119+
np.arange(0, 3, dtype="i4"),
120+
np.arange(0, 3, dtype="i4"),
121+
np.ones(3),
122+
],
123+
"shape": (10, pdict["M"] * 3),
124+
}
125+
jac["velocity"] = {
126+
"coo": [
127+
np.arange(3, 6, dtype="i4"),
128+
np.arange(0, 3, dtype="i4"),
129+
np.ones(3),
130+
],
131+
"shape": (10, pdict["M"] * 3),
132+
}
133+
jac["quaternion"] = {
134+
"coo": [
135+
np.arange(6, 10, dtype="i4"),
136+
np.arange(0, 4, dtype="i4"),
137+
np.ones(4),
138+
],
139+
"shape": (10, pdict["M"] * 4),
140+
}
131141

132142
else:
133-
rowcol_mass = ([0], [0])
134-
rowcol_pos = (range(1, 4), range(0, 3))
135-
rowcol_vel = (range(4, 7), range(0, 3))
136-
rowcol_quat = (range(7, 11), range(0, 4))
137-
138-
jac["mass"] = sparse.coo_matrix(([1.0], rowcol_mass), shape=(11, pdict["M"]))
139-
jac["position"] = sparse.coo_matrix(
140-
([1.0] * 3, rowcol_pos), shape=(11, pdict["M"] * 3)
141-
)
142-
jac["velocity"] = sparse.coo_matrix(
143-
([1.0] * 3, rowcol_vel), shape=(11, pdict["M"] * 3)
144-
)
145-
jac["quaternion"] = sparse.coo_matrix(
146-
([1.0] * 4, rowcol_quat), shape=(11, pdict["M"] * 4)
147-
)
143+
jac["mass"] = {
144+
"coo": [np.zeros(1, dtype="i4"), np.zeros(1, dtype="i4"), np.ones(1)],
145+
"shape": (11, pdict["M"]),
146+
}
147+
jac["position"] = {
148+
"coo": [
149+
np.arange(1, 4, dtype="i4"),
150+
np.arange(0, 3, dtype="i4"),
151+
np.ones(3),
152+
],
153+
"shape": (11, pdict["M"] * 3),
154+
}
155+
jac["velocity"] = {
156+
"coo": [
157+
np.arange(4, 7, dtype="i4"),
158+
np.arange(0, 3, dtype="i4"),
159+
np.ones(3),
160+
],
161+
"shape": (11, pdict["M"] * 3),
162+
}
163+
jac["quaternion"] = {
164+
"coo": [
165+
np.arange(7, 11, dtype="i4"),
166+
np.arange(0, 4, dtype="i4"),
167+
np.ones(4),
168+
],
169+
"shape": (11, pdict["M"] * 4),
170+
}
148171

149172
return jac
150173

@@ -191,7 +214,10 @@ def equality_jac_time(xdict, pdict, unitdict, condition):
191214
col.extend([i, i_ref])
192215
iRow += 1
193216

194-
jac["t"] = sparse.coo_matrix((data, (row, col)), shape=[iRow, len(xdict["t"])])
217+
jac["t"] = {
218+
"coo": [np.array(row, dtype="i4"), np.array(col, dtype="i4"), np.array(data)],
219+
"shape": (iRow, len(xdict["t"])),
220+
}
195221

196222
return jac
197223

@@ -239,26 +265,45 @@ def equality_jac_dynamics_mass(xdict, pdict, unitdict, condition):
239265
unit_t = unitdict["t"]
240266
num_sections = pdict["num_sections"]
241267

242-
jac["mass"] = sparse.lil_matrix((pdict["N"], pdict["M"]))
243-
jac["t"] = sparse.lil_matrix((pdict["N"], num_sections + 1))
268+
jac["mass"] = {"coo": [[], [], []], "shape": (pdict["N"], pdict["M"])}
269+
jac["t"] = {"coo": [[], [], []], "shape": (pdict["N"], num_sections + 1)}
244270

245271
for i in range(num_sections):
246272
a = pdict["ps_params"][i]["index_start"]
247273
n = pdict["ps_params"][i]["nodes"]
248274
b = a + n
249275

250276
if pdict["params"][i]["engineOn"]:
251-
jac["mass"][a:b, a + i : b + i + 1] = pdict["ps_params"][i]["D"] # lh
252-
jac["t"][a:b, i] = (
253-
-pdict["params"][i]["massflow"] / unit_mass * unit_t / 2.0
277+
278+
jac["mass"]["coo"][0].extend(sum([[i] * (n + 1) for i in range(a, b)], []))
279+
jac["mass"]["coo"][1].extend(list(range(a + i, b + i + 1)) * (n))
280+
jac["mass"]["coo"][2].extend(
281+
pdict["ps_params"][i]["D"].ravel(order="C").tolist()
282+
)
283+
284+
jac["t"]["coo"][0].extend(list(range(a, b)))
285+
jac["t"]["coo"][1].extend([i] * n)
286+
jac["t"]["coo"][2].extend(
287+
[-pdict["params"][i]["massflow"] / unit_mass * unit_t / 2.0] * n
254288
) # rh(to)
255-
jac["t"][a:b, i + 1] = (
256-
pdict["params"][i]["massflow"] / unit_mass * unit_t / 2.0
289+
jac["t"]["coo"][0].extend(list(range(a, b)))
290+
jac["t"]["coo"][1].extend([i + 1] * n)
291+
jac["t"]["coo"][2].extend(
292+
[pdict["params"][i]["massflow"] / unit_mass * unit_t / 2.0] * n
257293
) # rh(tf)
258294

259295
else:
260-
jac["mass"][a:b, a + i] = -1.0
261-
jac["mass"][a:b, a + i + 1 : b + i + 1] = np.eye(n)
296+
jac["mass"]["coo"][0].extend(list(range(a, b)))
297+
jac["mass"]["coo"][1].extend([a + i] * n)
298+
jac["mass"]["coo"][2].extend([-1.0] * n)
299+
jac["mass"]["coo"][0].extend(list(range(a, b)))
300+
jac["mass"]["coo"][1].extend(list(range(a + i + 1, b + i + 1)))
301+
jac["mass"]["coo"][2].extend([1.0] * n)
302+
303+
for key in jac.keys():
304+
jac[key]["coo"][0] = np.array(jac[key]["coo"][0], dtype="i4")
305+
jac[key]["coo"][1] = np.array(jac[key]["coo"][1], dtype="i4")
306+
jac[key]["coo"][2] = np.array(jac[key]["coo"][2], dtype="f8")
262307

263308
return jac
264309

@@ -977,30 +1022,23 @@ def equality_jac_6DoF_LGR_terminal(xdict, pdict, unitdict, condition):
9771022
else:
9781023
nRow = 1
9791024

980-
jac_base = {
981-
"position": np.zeros((nRow, pdict["M"] * 3)),
982-
"velocity": np.zeros((nRow, pdict["M"] * 3)),
983-
}
1025+
nCol = pdict["M"] * 3
1026+
jac["position"] = {"coo": [[], [], []], "shape": (nRow, nCol)}
1027+
jac["velocity"] = {"coo": [[], [], []], "shape": (nRow, nCol)}
9841028

9851029
for key in ["position", "velocity"]:
9861030

987-
jac_base[key][:, -3:] = 1.0
988-
jac_base_coo = sparse.coo_matrix(jac_base[key])
989-
jac[key] = {
990-
"coo": [jac_base_coo.row, jac_base_coo.col, jac_base_coo.data],
991-
"shape": jac_base[key].shape,
992-
}
993-
994-
for j in range(-3, 0):
1031+
for j in range(nCol - 3, nCol):
9951032
xdict[key][j] += dx
9961033
f_p = equality_6DoF_LGR_terminal(xdict, pdict, unitdict, condition)
9971034
xdict[key][j] -= dx
998-
jac_base[key][:, j] = (f_p - f_center) / dx
1035+
jac[key]["coo"][0].extend(list(range(nRow)))
1036+
jac[key]["coo"][1].extend([j] * nRow)
1037+
jac[key]["coo"][2].extend(((f_p - f_center) / dx).tolist())
9991038

1000-
for i in range(len(jac[key]["coo"][0])):
1001-
jac[key]["coo"][2][i] = jac_base[key][
1002-
jac[key]["coo"][0][i], jac[key]["coo"][1][i]
1003-
]
1039+
jac[key]["coo"][0] = np.array(jac[key]["coo"][0], dtype="i4")
1040+
jac[key]["coo"][1] = np.array(jac[key]["coo"][1], dtype="i4")
1041+
jac[key]["coo"][2] = np.array(jac[key]["coo"][2], dtype="f8")
10041042

10051043
return jac
10061044

@@ -1080,9 +1118,9 @@ def equality_jac_6DoF_rate(xdict, pdict, unitdict, condition):
10801118

10811119
f_center = equality_6DoF_rate(xdict, pdict, unitdict, condition)
10821120
nRow = len(f_center)
1083-
jac["position"] = sparse.lil_matrix((nRow, pdict["M"] * 3))
1084-
jac["quaternion"] = sparse.lil_matrix((nRow, pdict["M"] * 4))
1085-
jac["u"] = sparse.lil_matrix((nRow, pdict["N"] * 3))
1121+
jac["position"] = {"coo": [[], [], []], "shape": (nRow, pdict["M"] * 3)}
1122+
jac["quaternion"] = {"coo": [[], [], []], "shape": (nRow, pdict["M"] * 4)}
1123+
jac["u"] = {"coo": [[], [], []], "shape": (nRow, pdict["N"] * 3)}
10861124

10871125
iRow = 0
10881126

@@ -1098,86 +1136,113 @@ def equality_jac_6DoF_rate(xdict, pdict, unitdict, condition):
10981136
att = pdict["params"][i]["attitude"]
10991137
# attitude hold : angular velocity is zero
11001138
if att in ["hold", "vertical"]:
1101-
jac["u"][iRow : iRow + n * 3, a * 3 : (a + n) * 3] = np.eye(n * 3)
1139+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n * 3)))
1140+
jac["u"]["coo"][1].extend(list(range(a * 3, (a + n) * 3)))
1141+
jac["u"]["coo"][2].extend([1.0] * (n * 3))
11021142
iRow += n * 3
11031143

11041144
# kick-turn : pitch rate constant, roll/yaw rate is zero
11051145
elif att == "kick-turn" or att == "pitch":
1106-
jac["u"][iRow : iRow + n, a * 3 : (a + n) * 3 : 3] = np.eye(n)
1146+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n)))
1147+
jac["u"]["coo"][1].extend(list(range(a * 3, (a + n) * 3, 3)))
1148+
jac["u"]["coo"][2].extend([1.0] * n)
11071149
iRow += n
1108-
jac["u"][iRow : iRow + n, a * 3 + 2 : (a + n) * 3 + 2 : 3] = np.eye(n)
1150+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n)))
1151+
jac["u"]["coo"][1].extend(list(range(a * 3 + 2, (a + n) * 3 + 2, 3)))
1152+
jac["u"]["coo"][2].extend([1.0] * n)
11091153
iRow += n
1110-
jac["u"][iRow : iRow + n - 1, a * 3 + 1] = -1.0
1111-
jac["u"][
1112-
iRow : iRow + n - 1, (a + 1) * 3 + 1 : (a + n) * 3 + 1 : 3
1113-
] = np.eye(n - 1)
1154+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n - 1)))
1155+
jac["u"]["coo"][1].extend([a * 3 + 1] * (n - 1))
1156+
jac["u"]["coo"][2].extend([-1.0] * (n - 1))
1157+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n - 1)))
1158+
jac["u"]["coo"][1].extend(list(range((a + 1) * 3 + 1, (a + n) * 3 + 1, 3)))
1159+
jac["u"]["coo"][2].extend([1.0] * (n - 1))
11141160
iRow += n - 1
11151161

11161162
# pitch-yaw : pitch/yaw constant, roll ANGLE is zero
11171163
elif att == "pitch-yaw":
1118-
jac["u"][iRow : iRow + n - 1, a * 3 + 1] = -1.0
1119-
jac["u"][
1120-
iRow : iRow + n - 1, (a + 1) * 3 + 1 : (a + n) * 3 + 1 : 3
1121-
] = np.eye(n - 1)
1164+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n - 1)))
1165+
jac["u"]["coo"][1].extend([a * 3 + 1] * (n - 1))
1166+
jac["u"]["coo"][2].extend([-1.0] * (n - 1))
1167+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n - 1)))
1168+
jac["u"]["coo"][1].extend(list(range((a + 1) * 3 + 1, (a + n) * 3 + 1, 3)))
1169+
jac["u"]["coo"][2].extend([1.0] * (n - 1))
11221170
iRow += n - 1
1123-
jac["u"][iRow : iRow + n - 1, a * 3 + 2] = -1.0
1124-
jac["u"][
1125-
iRow : iRow + n - 1, (a + 1) * 3 + 2 : (a + n) * 3 + 2 : 3
1126-
] = np.eye(n - 1)
1171+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n - 1)))
1172+
jac["u"]["coo"][1].extend([a * 3 + 2] * (n - 1))
1173+
jac["u"]["coo"][2].extend([-1.0] * (n - 1))
1174+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n - 1)))
1175+
jac["u"]["coo"][1].extend(list(range((a + 1) * 3 + 2, (a + n) * 3 + 2, 3)))
1176+
jac["u"]["coo"][2].extend([1.0] * (n - 1))
11271177
iRow += n - 1
11281178
for k in range(n):
11291179
f_c = yb_r_dot(pos_i_[k + 1] * unit_pos, quat_i_[k + 1])
11301180
for j in range(3):
11311181
pos_i_[k + 1, j] += dx
11321182
f_p = yb_r_dot(pos_i_[k + 1] * unit_pos, quat_i_[k + 1])
11331183
pos_i_[k + 1, j] -= dx
1134-
jac["position"][iRow + k, (a + i + 1 + k) * 3 + j] = (
1135-
f_p - f_c
1136-
) / dx
1184+
jac["position"]["coo"][0].append(iRow + k)
1185+
jac["position"]["coo"][1].append((a + i + 1 + k) * 3 + j)
1186+
jac["position"]["coo"][2].append((f_p - f_c) / dx)
11371187
for j in range(4):
11381188
quat_i_[k + 1, j] += dx
11391189
f_p = yb_r_dot(pos_i_[k + 1] * unit_pos, quat_i_[k + 1])
11401190
quat_i_[k + 1, j] -= dx
1141-
jac["quaternion"][iRow + k, (a + i + 1 + k) * 4 + j] = (
1142-
f_p - f_c
1143-
) / dx
1191+
jac["quaternion"]["coo"][0].append(iRow + k)
1192+
jac["quaternion"]["coo"][1].append((a + i + 1 + k) * 4 + j)
1193+
jac["quaternion"]["coo"][2].append((f_p - f_c) / dx)
11441194
iRow += n
11451195

11461196
# same-rate : pitch/yaw is the same as previous section, roll ANGLE is zero
11471197
elif att == "same-rate":
1148-
jac["u"][iRow : iRow + n, a * 3 - 2] = -1.0
1149-
jac["u"][iRow : iRow + n, a * 3 + 1 : (a + n) * 3 + 1 : 3] = np.eye(n)
1198+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n)))
1199+
jac["u"]["coo"][1].extend([a * 3 - 2] * n)
1200+
jac["u"]["coo"][2].extend([-1.0] * n)
1201+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n)))
1202+
jac["u"]["coo"][1].extend(list(range(a * 3 + 1, (a + n) * 3 + 1, 3)))
1203+
jac["u"]["coo"][2].extend([1.0] * n)
11501204
iRow += n
1151-
jac["u"][iRow : iRow + n, a * 3 - 1] = -1.0
1152-
jac["u"][iRow : iRow + n, a * 3 + 2 : (a + n) * 3 + 2 : 3] = np.eye(n)
1205+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n)))
1206+
jac["u"]["coo"][1].extend([a * 3 - 1] * n)
1207+
jac["u"]["coo"][2].extend([-1.0] * n)
1208+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n)))
1209+
jac["u"]["coo"][1].extend(list(range(a * 3 + 2, (a + n) * 3 + 2, 3)))
1210+
jac["u"]["coo"][2].extend([1.0] * n)
11531211
iRow += n
11541212
for k in range(n):
11551213
f_c = yb_r_dot(pos_i_[k + 1] * unit_pos, quat_i_[k + 1])
11561214
for j in range(3):
11571215
pos_i_[k + 1, j] += dx
11581216
f_p = yb_r_dot(pos_i_[k + 1] * unit_pos, quat_i_[k + 1])
11591217
pos_i_[k + 1, j] -= dx
1160-
jac["position"][iRow + k, (a + i + 1 + k) * 3 + j] = (
1161-
f_p - f_c
1162-
) / dx
1218+
jac["position"]["coo"][0].append(iRow + k)
1219+
jac["position"]["coo"][1].append((a + i + 1 + k) * 3 + j)
1220+
jac["position"]["coo"][2].append((f_p - f_c) / dx)
11631221
for j in range(4):
11641222
quat_i_[k + 1, j] += dx
11651223
f_p = yb_r_dot(pos_i_[k + 1] * unit_pos, quat_i_[k + 1])
11661224
quat_i_[k + 1, j] -= dx
1167-
jac["quaternion"][iRow + k, (a + i + 1 + k) * 4 + j] = (
1168-
f_p - f_c
1169-
) / dx
1225+
jac["quaternion"]["coo"][0].append(iRow + k)
1226+
jac["quaternion"]["coo"][1].append((a + i + 1 + k) * 4 + j)
1227+
jac["quaternion"]["coo"][2].append((f_p - f_c) / dx)
11701228
iRow += n
11711229

11721230
# zero-lift-turn or free : roll hold
11731231
elif att == "zero-lift-turn" or att == "free":
1174-
jac["u"][iRow : iRow + n, a * 3 : (a + n) * 3 : 3] = np.eye(n)
1232+
jac["u"]["coo"][0].extend(list(range(iRow, iRow + n)))
1233+
jac["u"]["coo"][1].extend(list(range(a * 3, (a + n) * 3, 3)))
1234+
jac["u"]["coo"][2].extend([1.0] * n)
11751235
iRow += n
11761236

11771237
else:
11781238
print("ERROR: UNKNOWN ATTITUDE OPTION! ({})".format(att))
11791239
sys.exit()
11801240

1241+
for key in jac.keys():
1242+
jac[key]["coo"][0] = np.array(jac[key]["coo"][0], dtype="i4")
1243+
jac[key]["coo"][1] = np.array(jac[key]["coo"][1], dtype="i4")
1244+
jac[key]["coo"][2] = np.array(jac[key]["coo"][2], dtype="f8")
1245+
11811246
return jac
11821247

11831248

0 commit comments

Comments
 (0)