Crystal Lang: Kalman filter implementation

kalman.cr

require "matrix"
require "time"

class Array
  def middle
    if count < 3
      return self
    end

    from = (count / 4.0).round.to_i
    to = count - from

    if count.modulo(2)
      to -= 1
    end

    return self[from..to]
  end

  def median
    array = self.middle

    array.sum / array.size
  end
end

class Kalman
  property :state
  property :error_covariance

  getter :control

  getter :readings
  getter :readings_transpose

  getter :process_noise
  getter :sensor_noise
  getter :identity

  getter :model

  property :error
  property :system_uncertainty
  property :kalman_gain

  property :current_measurement
  property :current_timestep
  @current_timestep :: Float64

  property :update_time, :predict_time

  def initialize
     # x = Estimate value (intial state)
     @state = Matrix[
       [247.0], # Height,
       [0.0],   # Velocity
       [0.0]    # Acceleration
     ]

     # P = Initial uncertainty (process noise)
     @error_covariance = Matrix[
       [1000.0, 0.0, 0.0],  # We're not certian about height
       [0.0, 1000.0, 0.0],  # We're not certain about velocity
       [0.0, 0.0, 1000.0],  # We're not certain about acceleration
     ]

     # u = External motion (control signal)
     # There no control signal in our case, so it's just zeroes
     @control = Matrix[
       [0.0], # Height control signal, none
       [0.0], # Velocity control signal, none
       [0.0], # Acceleration control signal, none
     ]

     # H = Measurement function (measured params)
     # Constant
     @readings = Matrix[[
       1.0, # We can measure height
       0.0, # But we can't measure velocity
       0.0  # Not acceleration
     ]]

     @readings_transpose = @readings.transpose

     # Q = covariance of the process noise, and we know nothing about it :|
     # Constant
     @process_noise = Matrix[
       [0.0, 0.0, 0.0],
       [0.0, 0.0, 0.0],
       [0.0, 0.0, 0.0],
     ]

     # R = Measurement uncertainty (measurement noise)
     # Constant
     @sensor_noise = Matrix[
       [0.6] # We're sensing only height, and there some error associated with height
     ]

     # z = Current measurement (observed value)
     @current_measurement = Matrix[
       [0.0]
     ]

     # y = Error
     @error = Matrix[
       [0.0]
     ]

     # K = Kalman gain (current measurement uncertainty)
     @kalman_gain = Matrix[
       [0.0],
       [0.0],
       [0.0]
     ]

     # S = System Uncertainty
     @system_uncertainty = Matrix[
       [0.0]
     ]

     # I = Identity matrix
     @identity = Matrix.identity(3)

     @update_time = 0.0
     @predict_time = 0.0
     @current_timestep = 0.1

     # F = Next state function (process model)
     # This is where everything is goes nuts:
     # You need to come up with a matrix, that will represent a model
     # of changes between states. So if you apply this matrix to previous
     # state vector, you will get the next state prediction vector.
     @model = Matrix[
       [1.0, 1.0, -1.0],    # height  = height.prev + inflow - outflow
       [0.0, 1.0, 0.0],     # inflow  = inflow.prev
       [0.0, 0.0, 1.0]      # outflow = outflow.prev
     ]
   end

  def height
    self.state.first
  end

  def inflow
    self.state[1,0]
  end

  def outflow
    self.state[2,0]
  end

  def measure(value : Float64, timestep : Float64)
    self.current_measurement = Matrix[
      [value]
    ]

    self.current_timestep = timestep

    predict!
    update!
  end

  def predict!
    time = Time.now

    ## Prediction stage

    # x = (F * x) + B * u
    self.state = (self.model * self.state)
    # We have no control signal, so skip it for now
    # + self.control_model * self.control

    # P = F * P * Ft + Q
    self.error_covariance = self.model * self.error_covariance * self.model.transpose
    # We're have no idea about process noise, so skip it for now
    # + self.process_noise

    ## End

    self.predict_time = Time.now - time
  end

  def update!
    time = Time.now

    ## Update stage

    # y = z - (H * x)
    self.error = self.current_measurement - (self.readings * self.state)

    # S = H * P * Ht + R
    self.system_uncertainty = self.readings * self.error_covariance * self.readings_transpose + self.sensor_noise

    # K = P * Ht + S^-1
    self.kalman_gain = self.error_covariance * self.readings_transpose * system_uncertainty.inverse

    # x = x + (K * y)
    self.state = self.state + (self.kalman_gain * self.error)

    # P = (I - (K * H)) * P
    self.error_covariance = (self.identity - (self.kalman_gain * self.readings)) * self.error_covariance

    ## End

    self.update_time = Time.now - time
  end
end

class Parser
  getter :running
  getter :stats
  getter :files
  getter :buffer

  property :count

  getter :kalman

  getter :well_radius_pi
  getter :well_depth

  property :last_time
  property :last_value

  def initialize
    @files = ARGV.sort.uniq
    @running = false

    @stats = {
      lines: 0,
      empty: 0,
      comments: 0,
      bad_format: 0,
      bad_values: 0,
      dups: 0,
    }

    @count = 0

    @kalman = Kalman.new

    # Sensor position in the well
    @mount_level = 155.0
    @real_depth = 527.0
    @well_radius = 121.0
    @square_well_radius = @well_radius ** 2

    @well_depth = @real_depth - @mount_level
    @well_radius_pi = Math::PI * @square_well_radius

    @buffer = [] of Float64
  end

  def start!
    return if @running

    @running = true

    files.each do |file_path|
      break if !running

      read_file file_path
    end
  end

  def stop!
    @running = false
  end

  def read_file(path)
    date = parse_date(path)
    file = File.open(path)

    consume file, date
  ensure
    file.close if file
  end

  def consume(file, date)
    file.each_line do |line|
      self.stats[:lines] =+ 1
      data = extract_data(line, date)

      case data
      when :empty, :comment, :bad_format
        stats[data] =+ 1 if data.is_a? Symbol
        next
      else
        # This is just for Crystal
        if data.is_a? Tuple
          current_time, current_value = data

          # Collect 1 second worth of data, beacuse
          # feeding filter with 100ms steps are still
          # making a lot of low-freqeny noise
          if last_time == current_time
            self.buffer << current_value
          else
            if self.buffer.empty?
            else
              median_value = self.buffer.median
              self.buffer.clear

              filter current_time, median_value
            end
          end

          filter current_time, current_value

          self.last_time = current_time
        end
      end
    end
  end

  def filter(current_time : Time, current_value : Float64)
    self.count = self.count + 1

    # Oh boy, static languages.
    last_known_value = self.last_value
    last_known_time = self.last_time

    # Skip first step
    if last_known_time
      # When kalman is stabialized, we will filter out all the
      # outliers by simply dropping all the values that are diff
      # more than 5% than current filter estimation
      if count > 300
         if last_known_value
           max_delta = last_known_value * 0.05
           diff = (current_value - last_known_value).abs

           if diff > max_delta
             stats[:bad_value] =+ 1
             return
           end
         end
       end

      float_current_time = current_time.to_f
      float_last_time = last_known_time.to_f
      time_step = float_current_time - float_last_time

      if time_step == 0.0
        return
      end

      kalman.measure(current_value, time_step)
      new_value = kalman.height.round(2)

      display(new_value, current_value, current_time)

      self.last_value = new_value
    end
  end

  def display(new_value, current_value, current_time)
    return if count < 600
    return if last_value == new_value

    puts "%s, %.2f, %.2f" % [
      current_time.to_s("%F %T"),
      volume_of(current_value),
      volume_of(new_value),
    ]
  end

  def extract_data(line, date)
    # Oopsie, it breakes the crystal if done
    # at the one time. Should be fixed in HEAD
    string = line.gsub(/\0+/, "").gsub(/\s+/, "")

    if string.empty?
      return :empty
    end

    if string.starts_with? ';'
      return :comment
    end

    time_string, raw_value = string.split(",")

    if time_string.nil? || raw_value.nil?
      return :bad_format
    end

    time = parse_time(time_string, date)
    value = raw_value.to_f

    return time, value
  end

  def parse_time(string, date)
    if string.index '.'
      time, ms = string.split(".")
    else
      time = string
      ms = nil
    end

    # Reconstruct Time object
    return Time.parse("#{date}T#{time}Z", "%FT%TZ")
  rescue ex
    puts "Can't parse: '#{date}T#{time}Z': #{ex.message}"
    exit 1
  end

  def parse_date(path)
    match = path.match(/(\d{4}-\d{2}-\d{2}).csv/)

    if match
      return match[1]
    else
      puts "Can't get date from file path '#{path}'. Should contain YYYY-MM-DD.csv"
      exit 1
    end
  end

  def volume_of(obj : Float64)
    level = self.well_depth - obj

    return (self.well_radius_pi * level / 1000.0)
  end
end

parser = Parser.new
parser.start!