Linearise Nonlinear Model Automatically

%matplotlib notebook
from notebook_quick_setup import *

Beginning notebook setup...
        Added /home/jhewers/Documents/projects/jdrones/src to path
        Imported gymnasium version 0.27.1

pybullet build time: Feb  2 2023 13:13:41
/home/jhewers/.local/lib/python3.10/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.5
  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"

        Imported jdrones version unknown
        Imported scipy==1.7.3, numpy==1.23.5, pandas==1.3.5
        Imported functools, collections and itertools
        Imported tqdm (standard and trange)
        Imported seaborn==0.11.2, matplotlib==3.5.1
End of notebook setup

env = gymnasium.make('NonLinearDynamicModelDroneEnv-v0')

utrim = np.ones(4)*np.sqrt((env.model.mass*env.model.g)/(4*env.model.k_T))
xtrim = np.zeros(12)

nstates = len(xtrim)
ncontrols = len(utrim)
nv = nstates + ncontrols
stepSize = 0.01
minPertSize = 0.0001
cv = np.concatenate([xtrim, utrim])
pert = np.clip(np.abs(cv) * stepSize, minPertSize, np.inf)
dvar = 2.0 * pert
var = cv
varpert = np.zeros((2 * nv, nv))
for m in range(2 * nv):
    for n in range(nv):
        varpert[m, n] = var[n]
    if m > 1:
        ind = np.floor(0.5 * m).astype(int)
        if m % 2 == 1:
            varpert[m, ind] = var[ind] - pert[ind]
        else:
            varpert[m, ind] = var[ind] + pert[ind]

# evaluate the Jacobian
f = np.empty((nstates, 2 * nv))
for j in range(2 * nv):
    # Calculate the derivative corresponding to the perturbed state.
    xp = varpert[j, 0:nstates].T
    up = varpert[j, nstates:nv].T
    xd = env.calc_dstate(up, State.from_x(xp), env.model)
    xd = State(xd).to_x()
    # Now the function
    for i in range(nstates):
        f[i, j] = xd[i]
# calculate the Jacobian using numerical differentiation
J = np.empty((nstates, nv))
for m in range(nstates):
    for n in range(nv):
        a = f[m, 2 * n]
        b = f[m, 2 * n + 1]
        c = dvar[n]
        J[m, n]= (a - b) / c

sysMatrix = J[0:nstates, 0:nstates]
conMatrix = J[0:nstates, nstates:nv]

for row in sysMatrix:
    print(tuple(map(functools.partial(round, ndigits=2), row)))

(0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, -0.0, 0.0, 0.0, 0.0, 9.81, 0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0, -0.0, 0.0, -9.81, 0.0, 0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0, 0.0, -0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0)
(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0)
(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)

for row in conMatrix:
    print(tuple(map(functools.partial(round, ndigits=2), row)))

(0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0)
(0.84, 0.84, 0.84, 0.84)
(0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0)
(0.0, 0.0, 0.0, 0.0)
(0.0, -1.17, 0.0, 1.17)
(-1.17, 0.0, 1.17, 0.0)
(5.86, -5.86, 5.86, -5.86)