231 lines
6.5 KiB
Python
231 lines
6.5 KiB
Python
#!/usr/bin/env python
|
|
|
|
'''
|
|
fit best estimate of magnetometer offsets
|
|
'''
|
|
from __future__ import print_function
|
|
|
|
import sys, time, os, math
|
|
|
|
from argparse import ArgumentParser
|
|
parser = ArgumentParser(description=__doc__)
|
|
parser.add_argument("--no-timestamps", dest="notimestamps", action='store_true', help="Log doesn't have timestamps")
|
|
parser.add_argument("--condition", default=None, help="select packets by condition")
|
|
parser.add_argument("--noise", type=float, default=0, help="noise to add")
|
|
parser.add_argument("--mag2", action='store_true', help="use 2nd mag from DF log")
|
|
parser.add_argument("--radius", default=500.0, type=float, help="target radius")
|
|
parser.add_argument("--plot", action='store_true', help="plot points in 3D")
|
|
parser.add_argument("logs", metavar="LOG", nargs="+")
|
|
|
|
args = parser.parse_args()
|
|
|
|
from pymavlink import mavutil
|
|
from pymavlink.rotmat import Vector3
|
|
from pymavlink.rotmat import Matrix3
|
|
|
|
|
|
def noise():
|
|
'''a noise vector'''
|
|
from random import gauss
|
|
v = Vector3(gauss(0, 1), gauss(0, 1), gauss(0, 1))
|
|
v.normalize()
|
|
return v * args.noise
|
|
|
|
def select_data(data):
|
|
ret = []
|
|
counts = {}
|
|
for d in data:
|
|
mag = d
|
|
key = "%u:%u:%u" % (mag.x/20,mag.y/20,mag.z/20)
|
|
if key in counts:
|
|
counts[key] += 1
|
|
else:
|
|
counts[key] = 1
|
|
if counts[key] < 3:
|
|
ret.append(d)
|
|
print((len(data), len(ret)))
|
|
return ret
|
|
|
|
def constrain(v, min, max):
|
|
if v < min:
|
|
return min
|
|
if v > max:
|
|
return max
|
|
return v
|
|
|
|
def correct(mag, offsets, diag, offdiag):
|
|
'''correct a mag sample'''
|
|
diag.x = 1.0
|
|
mat = Matrix3(
|
|
Vector3(diag.x, offdiag.x, offdiag.y),
|
|
Vector3(offdiag.x, diag.y, offdiag.z),
|
|
Vector3(offdiag.y, offdiag.z, diag.z))
|
|
mag = mag + offsets
|
|
mag = mat * mag
|
|
return mag
|
|
|
|
def radius(mag, offsets, diag, offdiag):
|
|
'''return radius give data point and offsets'''
|
|
mag = correct(mag, offsets, diag, offdiag)
|
|
return mag.length()
|
|
|
|
def radius_cmp(a, b, offsets, diag, offdiag):
|
|
'''return +1 or -1 for for sorting'''
|
|
diff = radius(a, offsets, diag, offdiag) - radius(b, offsets, diag, offdiag)
|
|
if diff > 0:
|
|
return 1
|
|
if diff < 0:
|
|
return -1
|
|
return 0
|
|
|
|
def sphere_error(p, data):
|
|
x,y,z,r,dx,dy,dz,odx,ody,odz = p
|
|
if args.radius is not None:
|
|
r = args.radius
|
|
ofs = Vector3(x,y,z)
|
|
diag = Vector3(dx, dy, dz)
|
|
offdiag = Vector3(odx, ody, odz)
|
|
|
|
ret = []
|
|
for d in data:
|
|
mag = correct(d, ofs, diag, offdiag)
|
|
err = r - mag.length()
|
|
ret.append(err)
|
|
return ret
|
|
|
|
def fit_data(data):
|
|
from scipy import optimize
|
|
|
|
p0 = [0.0, 0.0, 0.0, # offsets
|
|
500.0, # radius
|
|
1.0, 1.0, 1.0, # diagonals
|
|
0.0, 0.0, 0.0 # offdiagonals
|
|
]
|
|
if args.radius is not None:
|
|
p0[3] = args.radius
|
|
p1, ier = optimize.leastsq(sphere_error, p0[:], args=(data))
|
|
if not ier in [1, 2, 3, 4]:
|
|
print(p1)
|
|
raise RuntimeError("Unable to find solution: %u" % ier)
|
|
if args.radius is not None:
|
|
r = args.radius
|
|
else:
|
|
r = p1[3]
|
|
return (Vector3(p1[0], p1[1], p1[2]),
|
|
r,
|
|
Vector3(p1[4], p1[5], p1[6]),
|
|
Vector3(p1[7], p1[8], p1[9]))
|
|
|
|
def magfit(logfile):
|
|
'''find best magnetometer offset fit to a log file'''
|
|
|
|
print("Processing log %s" % filename)
|
|
mlog = mavutil.mavlink_connection(filename, notimestamps=args.notimestamps)
|
|
|
|
data = []
|
|
|
|
last_t = 0
|
|
offsets = None
|
|
|
|
# now gather all the data
|
|
while True:
|
|
m = mlog.recv_match(condition=args.condition)
|
|
if m is None:
|
|
break
|
|
if m.get_type() == "SENSOR_OFFSETS":
|
|
# update current offsets
|
|
offsets = Vector3(m.mag_ofs_x, m.mag_ofs_y, m.mag_ofs_z)
|
|
if m.get_type() == "RAW_IMU":
|
|
mag = Vector3(m.xmag, m.ymag, m.zmag)
|
|
# add data point after subtracting the current offsets
|
|
if offsets is not None:
|
|
data.append(mag - offsets + noise())
|
|
if m.get_type() == "MAG" and not args.mag2:
|
|
offsets = Vector3(m.OfsX,m.OfsY,m.OfsZ)
|
|
mag = Vector3(m.MagX,m.MagY,m.MagZ)
|
|
data.append(mag - offsets + noise())
|
|
if m.get_type() == "MAG2" and args.mag2:
|
|
offsets = Vector3(m.OfsX,m.OfsY,m.OfsZ)
|
|
mag = Vector3(m.MagX,m.MagY,m.MagZ)
|
|
data.append(mag - offsets + noise())
|
|
|
|
print("Extracted %u data points" % len(data))
|
|
print("Current offsets: %s" % offsets)
|
|
|
|
orig_data = data
|
|
|
|
# find average values
|
|
avg = Vector3()
|
|
count = 0
|
|
for d in data:
|
|
avg += d
|
|
count += 1
|
|
avg /= count
|
|
|
|
# subtract average
|
|
data = []
|
|
for d in orig_data:
|
|
data.append(d - avg)
|
|
print("Average %s" % avg)
|
|
|
|
#data = select_data(data)
|
|
|
|
diag = Vector3(1,1,1)
|
|
offdiag = Vector3(0,0,0)
|
|
|
|
# remove initial outliers
|
|
if False:
|
|
data.sort(lambda a,b : radius_cmp(a,b,offsets,diag,offdiag))
|
|
data = data[len(data)/16:-len(data)/16]
|
|
|
|
# do an initial fit
|
|
(offsets, field_strength, diag, offdiag) = fit_data(data)
|
|
|
|
for count in range(3):
|
|
# sort the data by the radius
|
|
data.sort(lambda a,b : radius_cmp(a,b,offsets,diag,offdiag))
|
|
|
|
print("Fit %u : %s %s %s field_strength=%6.1f to %6.1f" % (
|
|
count, offsets, diag, offdiag,
|
|
radius(data[0], offsets,diag,offdiag), radius(data[-1], offsets,diag,offdiag)))
|
|
|
|
# discard outliers, keep the middle
|
|
data = data[len(data)/32:-len(data)/32]
|
|
|
|
# fit again
|
|
(offsets, field_strength, diag, offdiag) = fit_data(data)
|
|
|
|
print("Final : %s %s %s field_strength=%6.1f to %6.1f" % (
|
|
offsets, diag, offdiag,
|
|
radius(data[0], offsets, diag, offdiag), radius(data[-1], offsets, diag, offdiag)))
|
|
|
|
offsets -= avg
|
|
print("With average : %s" % offsets)
|
|
|
|
if args.plot:
|
|
data2 = [correct(d,offsets,diag,offdiag) for d in orig_data]
|
|
plot_data(orig_data, data2)
|
|
|
|
def plot_data(orig_data, data):
|
|
'''plot data in 3D'''
|
|
from mpl_toolkits.mplot3d import Axes3D
|
|
import matplotlib.pyplot as plt
|
|
|
|
fig = plt.figure()
|
|
for dd, c, p in [(orig_data, 'r', 1), (data, 'b', 2)]:
|
|
ax = fig.add_subplot(1, 2, p, projection='3d')
|
|
|
|
xs = [ d.x for d in dd ]
|
|
ys = [ d.y for d in dd ]
|
|
zs = [ d.z for d in dd ]
|
|
ax.scatter(xs, ys, zs, c=c, marker='o')
|
|
|
|
ax.set_xlabel('X')
|
|
ax.set_ylabel('Y')
|
|
ax.set_zlabel('Z')
|
|
plt.show()
|
|
|
|
total = 0.0
|
|
for filename in args.logs:
|
|
magfit(filename)
|