Coverage for src/amisc/examples/models.py: 95%
165 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-08-29 21:38 +0000
1"""Providing some example test functions"""
2import pickle
3import uuid
4from pathlib import Path
6import numpy as np
8from amisc.rv import NormalRV, UniformRV
9from amisc.system import ComponentSpec, SystemSurrogate
12def tanh_func(x, *args, A=2, L=1, frac=4, **kwargs):
13 """Simple tunable tanh function"""
14 return {'y': A*np.tanh(2/(L/frac)*(x-L/2)) + A + 1}
17def ishigami(x, *args, a=7.0, b=0.1, **kwargs):
18 """For testing Sobol indices: https://doi.org/10.1109/ISUMA.1990.151285"""
19 return {'y': np.sin(x[..., 0:1]) + a*np.sin(x[..., 1:2])**2 + b*(x[..., 2:3]**4)*np.sin(x[..., 0:1])}
22def borehole_func(x, *args, **kwargs):
23 """Model found at https://www.sfu.ca/~ssurjano/borehole.html
24 :returns vdot: Water flow rate in m^3/yr
25 """
26 rw = x[..., 0] # Radius of borehole (m)
27 r = x[..., 1] # Radius of influence (m)
28 Tu = x[..., 2] # Transmissivity (m^2/yr)
29 Hu = x[..., 3] # Potentiometric head (m)
30 Tl = x[..., 4] # Transmissivity (m^2/yr)
31 Hl = x[..., 5] # Potentiometric head (m)
32 L = x[..., 6] # Length of borehole (m)
33 Kw = x[..., 7] # Hydraulic conductivity (m/yr)
34 vdot = 2*np.pi*Tu*(Hu-Hl) / (np.log(r/rw) * (1 + (2*L*Tu/(np.log(r/rw)*Kw*rw**2)) + (Tu/Tl)))
36 return {'y': vdot[..., np.newaxis]}
39def wing_weight_func(x, *args, **kwargs):
40 """Model found at https://www.sfu.ca/~ssurjano/wingweight.html
41 :returns Wwing: the weight of the airplane wing (lb)
42 """
43 Sw = x[..., 0] # Wing area (ft^2)
44 Wfw = x[..., 1] # Weight of fuel (lb)
45 A = x[..., 2] # Aspect ratio
46 Lambda = x[..., 3] # Quarter-chord sweep (deg)
47 q = x[..., 4] # Dynamic pressure (lb/ft^2)
48 lamb = x[..., 5] # taper ratio
49 tc = x[..., 6] # Aerofoil thickness to chord ratio
50 Nz = x[..., 7] # Ultimate load factor
51 Wdg = x[..., 8] # Design gross weight (lb)
52 Wp = x[..., 9] # Paint weight (lb/ft^2)
53 Lambda = Lambda*(np.pi/180)
54 Wwing = 0.036*(Sw**0.758)*(Wfw**0.0035)*((A/(np.cos(Lambda))**2)**0.6)*(q**0.006)*(lamb**0.04)*\
55 (100*tc/np.cos(Lambda))**(-0.3)*((Nz*Wdg)**0.49) + Sw*Wp
57 return {'y': Wwing[..., np.newaxis]}
60def nonlinear_wave(x, *args, env_var=0.1**2, wavelength=0.5, wave_amp=0.1, tanh_amp=0.5, L=1, t=0.25, **kwargs):
61 """Custom nonlinear model of a traveling Gaussian wave for testing.
63 :param x: `(..., x_dim)`, input locations
64 :param env_var: variance of Gaussian envelope
65 :param wavelength: sinusoidal perturbation wavelength
66 :param wave_amp: amplitude of perturbation
67 :param tanh_amp: amplitude of tanh(x)
68 :param L: domain length of underlying tanh function
69 :param t: transition length of tanh function (as fraction of L)
70 :returns: `(..., y_dim)`, model output
71 """
72 # Traveling sinusoid with moving Gaussian envelope (theta is x2)
73 env_range = [0.2, 0.6]
74 mu = env_range[0] + x[..., 1] * (env_range[1] - env_range[0])
75 theta_env = 1 / (np.sqrt(2 * np.pi * env_var)) * np.exp(-0.5 * (x[..., 0] - mu) ** 2 / env_var)
76 ftheta = wave_amp * np.sin((2*np.pi/wavelength) * x[..., 1]) * theta_env
78 # Underlying tanh dependence on x1
79 fd = tanh_amp * np.tanh(2/(L*t)*(x[..., 0] - L/2)) + tanh_amp
81 # Compute model = f(theta, d) + f(d)
82 y = np.expand_dims(ftheta + fd, axis=-1) # (..., 1)
84 return {'y': y}
87def fire_sat_system(save_dir=None):
88 """Fire satellite detection system model from Chaudhuri 2018.
90 !!! Note "Some modifications"
91 Orbit will save outputs all the time; Power won't because it is part of an FPI loop. Orbit and Power can
92 return `np.nan` some of the time (optional, to test robustness of surrogate). Power and Attitude have an
93 `alpha` fidelity index that controls accuracy of the returned values. `alpha=0,1,2` corresponds to `8,4,0`
94 percent error.
96 :param save_dir: where to save model outputs
97 :returns: a `SystemSurrogate` object for the fire sat MD example system
98 """
99 Re = 6378140 # Radius of Earth (m)
100 mu = 3.986e14 # Gravitational parameter (m^3 s^-2)
101 eta = 0.22 # Power efficiency
102 Id = 0.77 # Inherent degradation of the array
103 thetai = 0 # Sun incidence angle
104 LT = 15 # Spacecraft lifetime (years)
105 eps = 0.0375 # Power production degradation (%/year)
106 rlw = 3 # Length to width ratio
107 nsa = 3 # Number of solar arrays
108 rho = 700 # Mass density of arrays (kg/m^3)
109 t = 0.005 # Thickness (m)
110 D = 2 # Distance between panels (m)
111 IbodyX = 6200 # kg*m^2
112 IbodyY = 6200 # kg*m^2
113 IbodyZ = 4700 # kg*m^2
114 dt_slew = 760 # s
115 theta = 15 # Deviation of moment axis from vertical (deg)
116 As = 13.85 # Area reflecting radiation (m^2)
117 c = 2.9979e8 # Speed of light (m/s)
118 M = 7.96e15 # Magnetic moment of earth (A*m^2)
119 Rd = 5 # Residual dipole of spacecraft (A*m^2)
120 rhoa=5.148e-11 # Atmospheric density (kg/m^3) -- typo in Chaudhuri 2018 has this as 1e11 instead
121 A = 13.85 # Cross-section in flight (m^2)
122 Phold = 20 # Holding power (W)
123 omega = 6000 # Max vel of wheel (rpm)
124 nrw = 3 # Number of reaction wheels
126 def orbit_fun(x, output_dir=None, pct_failure=0):
127 H = x[..., 0:1] # Altitude (m)
128 phi = x[..., 1:2] # Target diameter (m)
129 vel = np.sqrt(mu / (Re + H)) # Satellite velocity (m/s)
130 dt_orbit = 2*np.pi*(Re + H) / vel # Orbit period (s)
131 dt_eclipse = (dt_orbit/np.pi)*np.arcsin(Re / (Re + H)) # Eclipse period (s)
132 theta_slew = np.arctan(np.sin(phi / Re) / (1 - np.cos(phi / Re) + H/Re)) # Max slew angle (rad)
133 if np.random.rand() < pct_failure:
134 i = tuple([np.random.randint(0, N) for N in x.shape[:-1]])
135 i2 = tuple([np.random.randint(0, N) for N in x.shape[:-1]])
136 vel[i + (0,)] = np.nan
137 theta_slew[i2 + (0,)] = np.nan
138 y = np.concatenate((vel, dt_orbit, dt_eclipse, theta_slew), axis=-1)
139 if output_dir is not None:
140 files = []
141 id = str(uuid.uuid4())
142 for index in np.ndindex(*x.shape[:-1]):
143 fname = f'{id}_{index}.pkl'
144 with open(Path(output_dir) / fname, 'wb') as fd:
145 pickle.dump({'y': y[index + (slice(None),)]}, fd)
146 files.append(fname)
147 return {'y': y, 'files': files}
148 else:
149 return {'y': y}
151 def power_fun(x, alpha=(0,), *, output_dir=None, pct_failure=0):
152 pct = 1 - (2 - alpha[0]) * 0.04 if len(alpha) == 1 else 1 # extra pct error term
153 Po = x[..., 0:1] # Other power sources (W)
154 Fs = x[..., 1:2] # Solar flux (W/m^2)
155 dt_orbit = x[..., 2:3] # Orbit period (s)
156 dt_eclipse = x[..., 3:4] # Eclipse period (s)
157 Pacs = x[..., 4:5] # Power from attitude control system (W)
158 Ptot = Po + Pacs
159 Pe = Ptot
160 Pd = Ptot
161 Xe = 0.6 # These are power efficiencies in eclipse and daylight
162 Xd = 0.8 # See Ch. 11 of Wertz 1999 SMAD
163 Te = dt_eclipse
164 Td = dt_orbit - Te
165 Psa = ((Pe*Te/Xe) + (Pd*Td/Xd)) / Td
166 Pbol = eta*Fs*Id*np.cos(thetai)
167 Peol = Pbol * (1 - eps)**LT
168 Asa = Psa / Peol # Total solar array size (m^2)
169 L = np.sqrt(Asa*rlw/nsa)
170 W = np.sqrt(Asa/(rlw*nsa))
171 msa = 2*rho*L*W*t # Mass of solar array
172 Ix = msa*((1/12)*(L**2 + t**2) + (D+L/2)**2) + IbodyX
173 Iy = (msa/12)*(L**2 + W**2) + IbodyY # typo in Zaman 2013 has this as H**2 instead of L**2
174 Iz = msa*((1/12)*(L**2 + W**2) + (D + L/2)**2) + IbodyZ
175 Itot = np.concatenate((Ix, Iy, Iz), axis=-1)
176 Imin = np.min(Itot, axis=-1, keepdims=True)
177 Imax = np.max(Itot, axis=-1, keepdims=True)
178 if np.random.rand() < pct_failure:
179 i = tuple([np.random.randint(0, N) for N in x.shape[:-1]])
180 i2 = tuple([np.random.randint(0, N) for N in x.shape[:-1]])
181 Imin[i2 + (0,)] = np.nan
182 Asa[i + (0,)] = np.nan
183 y = np.concatenate((Imin, Imax*pct, Ptot*pct, Asa), axis=-1)
185 if output_dir is not None:
186 files = []
187 id = str(uuid.uuid4())
188 for index in np.ndindex(*x.shape[:-1]):
189 fname = f'{id}_{index}.pkl'
190 with open(Path(output_dir) / fname, 'wb') as fd:
191 pickle.dump({'y': y[index + (slice(None),)]}, fd)
192 files.append(fname)
193 return {'y': y, 'files': files}
194 else:
195 return {'y': y}
197 def attitude_fun(x, /, alpha=(0,)):
198 pct = 1 - (2 - alpha[0])*0.04 if len(alpha) == 1 else 1 # extra model fidelity pct error term
199 H = x[..., 0:1] # Altitude (m)
200 Fs = x[..., 1:2] # Solar flux
201 Lsp = x[..., 2:3] # Moment arm for solar radiation pressure
202 q = x[..., 3:4] # Reflectance factor
203 La = x[..., 4:5] # Moment arm for aerodynamic drag
204 Cd = x[..., 5:6] # Drag coefficient
205 vel = x[..., 6:7] # Satellite velocity
206 theta_slew = x[..., 7:8] # Max slew angle
207 Imin = x[..., 8:9] # Minimum moment of inertia
208 Imax = x[..., 9:10] # Maximum moment of inertia
209 tau_slew = 4*theta_slew*Imax / dt_slew**2
210 tau_g = 3*mu*np.abs(Imax - Imin)*np.sin(2*theta*(np.pi/180)) / (2*(Re+H)**3)
211 tau_sp = Lsp*Fs*As*(1+q)*np.cos(thetai) / c
212 tau_m = 2*M*Rd / (Re + H)**3
213 tau_a = (1/2)*La*rhoa*Cd*A*vel**2
214 tau_dist = np.sqrt(tau_g**2 + tau_sp**2 + tau_m**2 + tau_a**2)
215 tau_tot = np.max(np.concatenate((tau_slew, tau_dist), axis=-1), axis=-1, keepdims=True)
216 Pacs = tau_tot*(omega*(2*np.pi/60)) + nrw*Phold
217 y = np.concatenate((Pacs, tau_tot), axis=-1) * pct
218 return {'y': y}
220 orbit = ComponentSpec(orbit_fun, name='Orbit', exo_in=[0, 1], coupling_out=[0, 1, 2, 3], max_beta=(3, 3),
221 model_kwargs={'pct_failure': 0}, save_output=True)
222 power = ComponentSpec(power_fun, name='Power', truth_alpha=(2,), exo_in=[2, 3], max_alpha=(2,), max_beta=(3,)*5,
223 coupling_in={'Orbit': [1, 2], 'Attitude': [0]}, coupling_out=[4, 5, 6, 7], save_output=True,
224 model_kwargs={'pct_failure': 0})
225 attitude = ComponentSpec(attitude_fun, name='Attitude', truth_alpha=2, max_alpha=2, max_beta=(3,)*10,
226 exo_in=[0, 3, 4, 5, 6, 7], coupling_in={'Orbit': [0, 3], 'Power': [0, 1]},
227 coupling_out=[8, 9])
229 exo_vars = [NormalRV(18e6, 1e6, 'H'), NormalRV(235e3, 10e3, '\u03D5'), NormalRV(1000, 50, 'Po'),
230 NormalRV(1400, 20, 'Fs'), NormalRV(2, 0.4, 'Lsp'), NormalRV(0.5, 0.1, 'q'),
231 NormalRV(2, 0.4, 'La'), NormalRV(1, 0.2, 'Cd')]
232 coupling_vars = [UniformRV(2000, 6000, 'Vsat'), UniformRV(20000, 60000, 'To'), UniformRV(1000, 5000, 'Te'),
233 UniformRV(0, 4, 'Slew'), UniformRV(0, 12000, 'Imin'), UniformRV(0, 12000, 'Imax'),
234 UniformRV(0, 10000, 'Ptot'), UniformRV(0, 50, 'Asa'), UniformRV(0, 100, 'Pat'),
235 UniformRV(0, 5, 'tau_tot')]
236 surr = SystemSurrogate([orbit, power, attitude], exo_vars, coupling_vars, est_bds=500, save_dir=save_dir)
238 return surr