Kalman Filtered Nunchuck & Wii Motion Plus

Here's my latest. It's seems to be working well for me, but I have more cleanup and optimizations, as well as some more thought on correlating the linear axes from the nunchuk to the rotational axes of the WMP.

// Modified Kalman code using Roll, Pitch, and Yaw from a Wii MotionPlus and X, Y, and Z accelerometers from a Nunchuck.
// Also uses "new" style Init to provide unencrypted data from the Nunchuck to avoid the XOR on each byte.
// Requires the WM+ and Nunchuck to be connected to Arduino pins D3/D4 using the schematic from Knuckles904 and johnnyonthespot
// See http://randomhacksofboredom.blogspot.com/2009/07/motion-plus-and-nunchuck-together-on.html
// Kalman Code By Tom Pycke. http://tom.pycke.be/mav/71/kalman-filtering-of-imu-data
// Original zipped source with very instructive comments: http://tom.pycke.be/file_download/4
//
// Some of this code came via contributions or influences by Ed Simmons, Jordi Munoz, and Adrian Carter
//
// Created by Duckhead   v0.2   some cleanup and other optimization work remains.  also need to investigate correlation of nunchuk and WMP (which axis is which)

#include <math.h>
#include <Wire.h>
/////////////////////////////////////////

struct GyroKalman
{
  /* These variables represent our state matrix x */
  float x_angle,
        x_bias;

  /* Our error covariance matrix */
  float P_00,
        P_01,
        P_10,
        P_11;      

  /* 
   * Q is a 2x2 matrix of the covariance. Because we
   * assume the gyro and accelerometer noise to be independent
   * of each other, the covariances on the / diagonal are 0.
   *
   * Covariance Q, the process noise, from the assumption
   *    x = F x + B u + w
   * with w having a normal distribution with covariance Q.
   * (covariance = E[ (X - E[X])*(X - E[X])' ]
   * We assume is linear with dt
   */
  float Q_angle, Q_gyro;

  /*
   * Covariance R, our observation noise (from the accelerometer)
   * Also assumed to be linair with dt
   */
  float R_angle;
};


struct GyroKalman rollData;
struct GyroKalman pitchData;
struct GyroKalman yawData;

// test for 3 axis
int readingsX;
int readingsY;
int readingsZ;                // the readings
// End of 3 axis stuff
int index = 0;                            // the index of the current reading
int inputPin =0;                          // Gyro Analog input
//WM+ stuff
byte data[6]; //six data bytes
int yaw = 0;
int pitch = 0;
int roll = 0; //three axes
int yaw0 = 0;
int pitch0 = 0;
int roll0 = 0; //calibration zeroes
bool slowYaw;
bool slowPitch;
bool slowRoll;

const int wmpSlowToDegreePerSec = 20;
const int wmpFastToDegreePerSec = 4;

///////////////////////////////////////

float accelAngleX=0;      //NunChuck X angle
float accelAngleY=0;      //NunChuck Y angle
float accelAngleZ=0;      //NunChuck Z angle
byte outbuf[6];            // array to store arduino output
int cnt = 0;            //Counter
unsigned long lastread=0;

/*
 * R represents the measurement covariance noise.  In this case,
 * it is a 1x1 matrix that says that we expect 0.3 rad jitter
 * from the accelerometer.
 */
static const float      R_angle      = .3; //.3 default

/*
 * Q is a 2x2 matrix that represents the process covariance noise.
 * In this case, it indicates how much we trust the acceleromter
 * relative to the gyros.
 */
static const float      Q_angle      = 0.01; //(Kalman)
static const float      Q_gyro      = 0.04; //(Kalman)

//These are the limits of the values I got out of the Nunchuk accelerometers (yours may vary).
const int lowX = -215;
const int highX = 221;

const int lowY = -215;
const int highY = 221;

const int lowZ = -215;
const int highZ = 255;

/*
 * Nunchuk accelerometer value to degree conversion.  Output is -90 to +90 (180 degrees of motion).  
 */
float angleInDegrees(int lo, int hi, int measured) {
  float x = (hi - lo)/180.0;
  return (float)measured/x;
}

void switchtonunchuck(){
/*  PORTE ^= 32; // Toggle D3 LOW
  PORTG |= 32; // Force D4 HIGH */ // Fast switching of pins... Arduino MEGA specific...
  digitalWrite(3, LOW);
  digitalWrite(4, LOW);
  digitalWrite(4, HIGH);

}

void switchtowmp(){
/*  PORTG ^= 32; // Toggle D4 LOW
  PORTE |= 32; // Force D3 HIGH */ // Fast switching of pins... Arduino MEGA Specific pin maps, so switched back to 'normal' code for example
  digitalWrite(3, LOW);
  digitalWrite(4, LOW);
  digitalWrite(3, HIGH);
}

void nunchuck_init () //(nunchuck)
{
  // Uses New style init - no longer encrypted so no need to XOR bytes later... might save some cycles
  Wire.beginTransmission (0x52);      // transmit to device 0x52
  Wire.send (0xF0);            // sends memory address
  Wire.send (0x55);            // sends sent a zero.
  Wire.endTransmission ();      // stop transmitting
  delay(100);
  Wire.beginTransmission (0x52);      // transmit to device 0x52
  Wire.send (0xFB);            // sends memory address
  Wire.send (0x00);            // sends sent a zero.
  Wire.endTransmission ();      // stop transmitting
}

void setup()
{
  Serial.begin(115200);
/*  DDRG |= 32; // Force D3 to Output // Arduino MEGA "fast" pin config. Reverted to 'normal' for example post.
  DDRE |= 32; // Force D4 to Output
  PORTG ^= 32; // Toggle D3 HIGH
  PORTE ^= 32; // Toggle D4 HIGH */

  pinMode(3, OUTPUT);
  pinMode(4, OUTPUT);
  digitalWrite(3, HIGH);
  digitalWrite(4, HIGH);

  Wire.begin ();            // join i2c bus with address 0x52 (nunchuck)
  switchtowmp();
  wmpOn();
  calibrateZeroes();
  switchtonunchuck();
  nunchuck_init (); // send the initilization handshake

  initGyroKalman(&rollData, Q_angle, Q_gyro, R_angle);
  initGyroKalman(&pitchData, Q_angle, Q_gyro, R_angle);
  initGyroKalman(&yawData, Q_angle, Q_gyro, R_angle);
  
  lastread = millis();
}

void initGyroKalman(struct GyroKalman *kalman, const float Q_angle, const float Q_gyro, const float R_angle) {
  kalman->Q_angle = Q_angle;
  kalman->Q_gyro  = Q_gyro;
  kalman->R_angle = R_angle;
  
  kalman->P_00 = 0;
  kalman->P_01 = 0;
  kalman->P_10 = 0;
  kalman->P_11 = 0;
}

void send_zero () //(nunchuck)
{
  Wire.beginTransmission (0x52);      // transmit to device 0x52 (nunchuck)
  Wire.send (0x00);            // sends one byte (nunchuck)
  Wire.endTransmission ();      // stop transmitting (nunchuck)
}

void wmpOn(){
  Wire.beginTransmission(0x53);//WM+ starts out deactivated at address 0x53
  Wire.send(0xfe); //send 0x04 to address 0xFE to activate WM+
  Wire.send(0x04);
  Wire.endTransmission(); //WM+ jumps to address 0x52 and is now active
}

void wmpOff(){
  Wire.beginTransmission(82);
  Wire.send(0xf0);//address then
  Wire.send(0x55);//command
  Wire.endTransmission();
}

void wmpSendZero(){
  Wire.beginTransmission(0x52);//now at address 0x52
  Wire.send(0x00); //send zero to signal we want info
  Wire.endTransmission();
}

void calibrateZeroes(){
  long y0 = 0;
  long p0 = 0;
  long r0 = 0;
  const int avg = 10;
  for (int i=0;i<avg; i++){
    wmpSendZero();
    Wire.requestFrom(0x52,6);
    for (int t=0;t<6;t++){
      data[t]=Wire.receive();
    }
    y0+=(((data[3] >> 2) << 8)+data[0]);
    r0+=(((data[4] >> 2) << 8)+data[1]);
    p0+=(((data[5] >> 2) << 8)+data[2]);
  }
  
  yaw0 = y0/avg;
  roll0 = r0/avg;
  pitch0 = p0/avg;
}

void receiveData(){
  receiveRaw();
  //see http://wiibrew.org/wiki/Wiimote/Extension_Controllers#Wii_Motion_Plus
  //for info on what each byte represents
  yaw -= yaw0;
  roll -= roll0;
  pitch -= -pitch0;
}

void receiveRaw(){
  wmpSendZero(); //send zero before each request (same as nunchuck)
  Wire.requestFrom(0x52,6);//request the six bytes from the WM+
  for (int i=0;i<6;i++){
    data[i]=Wire.receive();
  }
  yaw=((data[3] >> 2) << 8)+data[0];
  roll=((data[4] >> 2) << 8)+data[1];
  pitch=((data[5] >> 2) << 8)+data[2];
  
  slowPitch = data[3] & 1;
  slowYaw = data[3] & 2;
  slowRoll = data[4] & 2;
}

void loop()
{
  unsigned long now = millis();
  float dt = ((float)(now - lastread))/1000.0; //compute the time delta since last iteration, in seconds.

  if (dt >= 0.01) {  //Only process delta angles if at least 1/100 of a second has elapsed
    lastread = now;

    switchtowmp();
    receiveData();
    switchtonunchuck();
    readNunchuck();

    // Handle the fast/slow bits of the Wii MotionPlus
    const int rollScale  = slowRoll  ? wmpSlowToDegreePerSec : wmpFastToDegreePerSec;
    const int pitchScale = slowPitch ? wmpSlowToDegreePerSec : wmpFastToDegreePerSec;
    const int yawScale   = slowYaw   ? wmpSlowToDegreePerSec : wmpFastToDegreePerSec;
    
    readingsX = (roll/rollScale); // read from the gyro sensor 
    readingsY = (pitch/pitchScale); // read from the gyro sensor 
    readingsZ = (yaw/yawScale); // read from the gyro sensor 
    
    float tmpX = accelAngleX * (PI/180.0);
    float tmpY = accelAngleY * (PI/180.0);
    float tmpZ = accelAngleZ * (PI/180.0);

    predict(&rollData, readingsX, dt);
    predict(&pitchData, readingsY, dt);
    predict(&yawData, readingsZ, dt);
  
    float rollAngle = update(&rollData, tmpX);
    float pitchAngle = update(&pitchData, tmpY);
    float yawAngle = update(&yawData, tmpZ);

    Serial.print("Roll: ");
    Serial.print(rollAngle*(180.0/PI));
    Serial.print("Pitch: ");
    Serial.print(pitchAngle*(180.0/PI));
    Serial.print("Yaw: ");
    Serial.println(yawAngle*(180.0/PI));
  }
}

//Stitch Here

Duck