├── .gitignore ├── README.md ├── UNLICENSE.md ├── experiment_plots └── .gitignore ├── kalman.py ├── main.py ├── matplotlibrc ├── reading_plots └── .gitignore ├── requirements.txt ├── sensor.py ├── traincar.py └── utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.jpg 2 | *.pdf 3 | *.png 4 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | #Sensor Fusion Tutorial 2 | Or how to avoid the infamous poor car 🚂🚃🚃💩🚃 3 | 4 | ## What _is_ this sensor fusion thing? 5 | 6 | This blog post is about sensor fusion. You might think you don’t know what that means, but don’t worry, you do. It’s something you do all the time, as part of your daily life. In fact, a lot of it is done by your nervous system autonomously, so you might not even notice that it’s there unless you look for it, or if something goes wrong. 7 | 8 | Take, for example, riding a bicycle. If you know how to do it, you likely remember the frustrating days you spent getting on only to fall back down right away. In part, what made it so hard is that riding a bike requires several separate internal sensors to work in concert with one another, and at a high level of precision. Your brain must learn how to properly interpret and integrate visual cues with what your hands and feet perceive and with readings from your vestibular system (i.e., the “accelerometers” and “gyroscopes” in your inner ear). Perhaps going in a straight line is doable without visual cues, but I am willing to bet that taking proper turns while blindfolded is impossible for most people. That is because the visual cues add information about the state of the bicycle that none of the other sensors can provide such as what is the angle of the curve ahead, what angle am I currently turning at, and am I heading straight for a tree instead of the middle of the road? On the other hand, getting rid of the information from the vestibular system would be even more detrimental. The visual system, while very good at finding upcoming obstacles, is quite bad at determining small deviations from vertical, and it probably won’t know you are falling until you are already halfway down. 9 | 10 | Ok, so both sensors are needed for riding a bike successfully, but why fuse the information? Couldn’t your eyes make sure that you’re not about to hit an open door, while your inner ear makes sure you’re staying upright? Of course, they could, but it turns out that integrating the two sources of data goes a long way toward eliminating noise from the measurement. Even though the eyes are very bad at determining equilibrium, they still provide some useful information. Your brain could throw this information away, or it could spend a tiny bit of extra energy and use it to significantly improve the accuracy of the inner ear sensor. 11 | 12 | Engineers have been quick to catch onto this—at least soon after electronic sensors became a thing—and so have developed similar techniques to be used in control systems where the same variable can be measured in several different ways: the velocity of cars, the altitude of airplanes or the [attitude](https://en.wikipedia.org/wiki/Attitude_indicator) of space rockets. In the rest of the blog post I'll introduce some of these techniques while focusing on a relatively pedestrian end-goal: bettering my train commute. 13 | 14 | ## Sensing the number of people in a train car 15 | 16 | Here at Datascope most of us take the CTA (Chicago Transit Authority) to work and back. My particular commute is 45 minutes each way on the [Red Line](https://en.wikipedia.org/wiki/Red_Line_(CTA)), which is the most crowded line in the city. As a result, it is sometimes hard to find a seat, and, on busy days, even to get on the train. Sometimes, like on baseball game days, this is unavoidable, since there are more people in the station than there is space on the train. On most days, however, this is merely a problem of distribution: some cars might be full, while others are only half full. Without waiting for the whole train to pass by and looking in every one of 20 or so cars, there is no way to know. 17 | 18 | So, keeping this in mind, what could we do if we could measure or estimate the number of people in each car, in real time? For one, we could display a distribution of car occupancy, even before the train arrives. People could wait outside the emptiest car, and not try and crowd into an already crowded one. This would not just make the rides more comfortable, but also speed up boarding, which is a major factor of train delays, at least according to the [Chicago Transit Authority courtesy guide](http://www.transitchicago.com/courtesy/). 19 | 20 | In addition, the occupancy estimate could solve a personal pet peeve of mine: [the poop car](http://www.redeyechicago.com/redeye-avoiding-poop-winter-cta-20141117-story.html). This is a Chicago phenomenon that happens in the winter: you think you scored a nice, relatively empty car, at a busy time too! It looks almost too good to be true, and it turns out it is. The car is filled with some stench or another, making it nearly impossible to breathe, so much so that, at the next stop almost all riders get off and switch to other, more crowded cars, while a new set of “suckers” walks into the trap. Broad City illustrated this phenomenon admirably: 21 | 22 |
23 | Broad City smells a poo 24 |
25 | 26 | Now, if we had a good estimate of the number of people in each car, we could easily see that one car is much emptier than the rest of the train, while the neighboring cars are fuller. For the Red Line in winter, that is a certain sign of either unbearable stench or some other, equally obnoxious problem. 27 | 28 | ## Sensor models 29 | 30 | A simple and cheap way to get a (very rough) headcount in an enclosed space is to measure the level of carbon dioxide in the space. Since people breathe out more CO2 than they breathe in, a higher concentration of CO2 usually correlates with more people. How much, however, depends a lot on the space, (and maybe on other things, like light, time of day, or HVAC dynamics). There is no master function that says “when your sensor reading 1000ppm of CO2 there are 15 people in your room,” and no such function can be calculated for a general case. This is why we need to build our own function, and we’re going to do this by first running an experiment and using the data to build a sensor model. We install a cheap CO2 sensor in each train car, and lump all sources of error into generic “sensor noise”. In our experiment we take readings of both the CO2 level and the number of people present at that time in the train car, many times throughout a day. In code, it looks a little bit like this: 31 | 32 | ```python 33 | class Sensor(object): 34 | ... 35 | def read(self, variable): 36 | """Given `variable` occupants, returns a sensor reading""" 37 | ... 38 | 39 | def fit(self, data): 40 | """Fits a sensor model to the experimental `data`""" 41 | ... 42 | 43 | co2_sensor = Sensor( 44 | "CO_2", intersect=350, slope=15, sigma=10, 45 | round_level=500, proc_sigma=30, units="ppm") 46 | 47 | for _ in range(datapoints): 48 | occupants = random.randrange(max_occupants + 1) 49 | sensor_data.append((occupants, co2_sensor.read(occupants))) 50 | 51 | co2_sensor.fit(sensor_data) 52 | ``` 53 | 54 | Let’s say we get a figure like this: 55 | 56 | ![CO2 sensor model](https://datascopeanalytics.com/blog/sensor-fusion-tutorial/experiment_plots/co_2.svg) 57 | 58 | The above figure helps us build a sensor model for our particular car. A simple model fitting a line to the points in the figure tells us, for every given reading, how many people we expect to find in the room.[ref]Even though we use the model in a predictive fashion—given a reading, we look up the number of people that corresponds to that reading—we fit the line considering the number of occupants as the independent variable; doing it the other way around would generate a bad fit for our purposes[/ref] 59 | 60 | We could stop here, report that number of people, and call it done. But we’re no fools, we see there’s a LOT of noise in the reading. And so, we also calculate the standard error for our fit, and get a rough estimate of how good our model is. In the above figures, you can see the uncertainty of our sensor model as a shaded gradient. 61 | 62 | Let’s talk about what that means for a minute. Say we check our sensor and get a reading of 1092ppm (a number I obtained by randomly mashing my keyboard). We see that this corresponds, on average, to there being approximately 46 people in the train car, but it’s pretty much just as likely that there are 40 or 50 of them. The car could even be empty—that option is still within two standard deviations of the mean.[ref]If you think that this is a simplistic model, that’s because it is. This is more or less the smallest building block of sensor fusion theory, and like any “smallest block” it has its problems. For example, you might have noticed that there’s a non-zero probability there are negative occupants in the room. While it’s true that we can leverage this fact to get a better predictor, allowing negative values for people is not strictly detrimental, as we can think of that as “the reading was so low, it was even *more* unlikely there was anyone there”[/ref] 63 | 64 | ![Sensor Fusion](https://datascopeanalytics.com/blog/sensor-fusion-tutorial/reading_plots/1_co2.svg) 65 | 66 | ## Multiple readings between stations 67 | 68 | Now that we have a sensor model, let’s actually get some data! Let’s start looking at what’s happening in the car between two stations, something we can determine easily, either from an accelerometer on the car or from the great [train tracker app](http://www.transitchicago.com/traintracker/) CTA provides. The nice thing about the train between stations is that very few people change cars at that time (it is in fact illegal to do so). So, we can assume that the number of people in the car stays constant between measurements. We already have a sensor reading, so let’s grab one more. 69 | 70 | ![Sensor Fusion](https://datascopeanalytics.com/blog/sensor-fusion-tutorial/reading_plots/2_co2.svg) 71 | 72 | This one is centered somewhere around 46 people and has, unsurprisingly, the same standard deviation as the first reading—this is a side effect of our model. Now comes the fun part: we merge the two readings into one probability curve: 73 | 74 | ![Sensor Fusion](https://datascopeanalytics.com/blog/sensor-fusion-tutorial/reading_plots/3_co2.svg) 75 | 76 | Mathematically, we multiply the two blue distributions and normalize the result as needed. Because our distributions are Gaussian this multiplication can be done analytically: 77 | $$\mu = \frac{\sigma_{p}^2\mu_{up}+\sigma_{up}^2\mu_{p}}{\sigma_{up}^2+\sigma_{p}^2}, \quad \sigma=\frac{\sigma_{up}^2\sigma_{p}^2}{\sigma_{up}^2+\sigma_{p}^2}$$ 78 | In code, it looks something like this: 79 | 80 | ```python 81 | class Gaussian(object): 82 | ... 83 | def bayesian_update(self, other): 84 | sigma_sum = self.sigma**2 + other.sigma**2 85 | self.mu = ((self.sigma**2) * other.mu + (other.sigma**2) * self.mu) / sigma_sum 86 | self.sigma = np.sqrt(((self.sigma * other.sigma)**2) / sigma_sum) 87 | ``` 88 | 89 | ## Multiple sensors 90 | 91 | We can apply the process in the previous section to multiple sensors that measure the same thing. For example, let’s say we want to install a temperature sensor, whose measurements also (loosely) correlate with the number of people in the room. We install the new sensor: 92 | 93 | ```python 94 | temp_sensor = Sensor( 95 | "Temperature", intersect=0, slope=0.25, sigma=5, 96 | round_level=10, proc_sigma=5, units="$^{\\circ}$C") 97 | ``` 98 | 99 | After doing the required experiments, the model for the temperature sensor looks like this : 100 | 101 | ![Temperature.png](https://datascopeanalytics.com/blog/sensor-fusion-tutorial/experiment_plots/Temperature.svg) 102 | 103 | Now, each measurement between the two stations could come from one sensor, or the other. Since we know how much to trust each sensor, we use the same method as before to fuse them together. The math does not change, since we’re still multiplying Gaussians, and the same equations apply. 104 | 105 | The first temperature measurement, for example, would fuse with the other two measurements like this: 106 | 107 | ![Sensor Fusion](https://datascopeanalytics.com/blog/sensor-fusion-tutorial/reading_plots/4_co2.svg) 108 | 109 | A 5 minute train ride might look like this, where the red line shows the true occupancy of the car: 110 | 111 | 114 | 115 | To make things easier, I made some custom classes to abstract away measurements, and used them as following: 116 | 117 | ```python 118 | class Reading(Gaussian): 119 | 120 | def __init__(self, sensor, truth, timestamp=None): 121 | """Generate a reading from `sensor` at `timestamp`, 122 | given `truth` people""" 123 | ... 124 | 125 | class Estimate(Gaussian): 126 | ... 127 | 128 | def add_reading(self, reading): 129 | self.reading_vector.append(reading) 130 | self.update(reading) 131 | 132 | def update(self, reading): 133 | ... 134 | self.bayesian_update(reading) 135 | 136 | reading_1 = Reading(co2_sensor, 60) 137 | reading_2 = Reading(temp_sensor, 60) 138 | 139 | estimate = Estimate() 140 | estimate.add_reading(reading_1) 141 | estimate.add_reading(reading_2) 142 | 143 | reading_3 = Reading(co2_sensor, 60) 144 | estimate.add_reading(reading_3) 145 | ``` 146 | 147 | ## Moving targets 148 | 149 | Until now we assumed the number of people in the car stays constant. That’s a decent assumption between two stations, but if we consider a whole route, this is certainly not true. Fortunately, the previous approach can be easily adapted to moving targets—and by this I mean “changing number of people”, not moving people between stations. 150 | 151 | In order to estimate changing number of people we need to introduce the concept of time between measurements, and also build a model of human behavior between measurements[ref]in fact, we are already using a model for the human behavior between measurements: we model the number of people as non-changing[/ref]. Let’s assume that, at every station, a random number of people enter or leave, except at the last station, when everyone leaves the train. In mathematical terms this means that, every time we stop at a station, the error on our estimate of the number of people in the car will have to widen, and since we don’t know how much we’ll just set its sigma to infinity. In code, this is done by either creating a brand new `Estimate` or setting its `sigma` to `None`: 152 | ```python 153 | estimate=Estimate() 154 | # or 155 | estimate.sigma = None 156 | ``` 157 | However, as the train is moving between the stations we will be getting measurements again, and our estimation will narrow. Over several stations our estimate will look something like this: 158 | 159 | 162 | 163 | In the end we get a filtered version of the signal, translated from CO2 and temperature levels to number of people in the room. We also get a confidence interval at every point in time. Perhaps, more importantly, we obtained these estimates in *real time*. What this means is that we didn’t have to wait until the train ride was over, or until we got to the next station in order to average all the data and get a good estimate. Instead, we included every single measurement as it came in, and we could present the people waiting at the next station a constantly updating estimate of how many people were in each car. 164 | 165 | ## Better models 166 | 167 | What we did in the previous sections is a well know, and quite brilliant method, called *Kalman filtering*. It’s one of the technological advances that helped Apollo 11 get to the moon, and variants of it are used nowadays on pretty much any non-trivial control system. An old joke says that every time the late Rudolf Kálmán got on a plane he would state his name and ask to be shown “his filter”. A very similar approach, called *Kalman smoothing* can be taken if we do not need real time results, but instead can use the full data at once to determine the most likely outcome. In our case (due to the fact that we assume no movement between cars between stations ) a Kalman smoother would have resulted in the estimate right before each station being considered as the estimate for that whole leg of the trip. Because of its overall much better estimate, smoothing is usually preferred to filtering when real-time results are not necessary. The downside would be that we wouldn’t be able to tell which car is the smelly car until the end of the line, and then it’s no better than a quick in-person smell test of the whole train. 168 | 169 | Improvements on this approach are numerous, and I’m sure some of you are already thinking about what they might be. I’ll list only a few: 170 | 171 | - non-linear sensor models: we fit a polynomial or exponential curve to the experimental data 172 | - non-constant sensor noise model: the standard deviation is different at different sensor readings, not constant across the range of readings 173 | - time-varying sensor model: the way the sensors behave might change with time of day, or other factors (e.g. air conditioning) 174 | - more complex human dynamics: the way people enter and leave a car is a lot more complicated than we assumed; we could, for example, take into account that people are more likely to enter or leave during rush hour; people might prefer to enter into emptier cars; or they might be actively leaving the poop car 175 | - some other stuff I haven’t thought of, but you probably have: add it to the comments! 176 | 177 | ## See also 178 | 179 | I purposefully kept this post light, especially in terms of math. Just a quick look at the wikipedia page for Kalman filtering is enough to turn most people off the whole subject. I wanted to make sure that you first and foremost follow the concepts presented here intuitively, rather than mathematically. However, I bet there are some of you who want more “meat” than what I’ve given here. I hear you. Consider this an appetizer. I’ll let you do your own cooking for the main meal, but I’ll give a list of potential ingredients: 180 | 181 | - [the github repo](https://github.com/datascopeanalytics/sensor_fusion) with the full code backing this post 182 | - [a succint technical version of this presentation, for those familiar with complex statistics notation](http://www.stats.ox.ac.uk/~steffen/teaching/bs2HT9/kalman.pdf) 183 | - the Wikipedia page for [Kalman filtering](https://en.wikipedia.org/wiki/Kalman_filter) 184 | - [pykalman](https://pykalman.github.io/), a very simple python library I have used in the past 185 | - Steve LaValle’s relatively accessible [blog post](https://developer3.oculus.com/blog/sensor-fusion-keeping-it-simple/) on how they do sensor fusion for the Oculus while avoiding Kalman Filtering 186 | - a very nice and simple explanation of [particle filtering](https://www.youtube.com/watch?v=aUkBa1zMKv4), which replaces assumptions of Gaussian distributions with hard work (on the part of the computer) 187 | -------------------------------------------------------------------------------- /UNLICENSE.md: -------------------------------------------------------------------------------- 1 | This is free and unencumbered software released into the public domain. 2 | 3 | Anyone is free to copy, modify, publish, use, compile, sell, or 4 | distribute this software, either in source code form or as a compiled 5 | binary, for any purpose, commercial or non-commercial, and by any 6 | means. 7 | 8 | In jurisdictions that recognize copyright laws, the author or authors 9 | of this software dedicate any and all copyright interest in the 10 | software to the public domain. We make this dedication for the benefit 11 | of the public at large and to the detriment of our heirs and 12 | successors. We intend this dedication to be an overt act of 13 | relinquishment in perpetuity of all present and future rights to this 14 | software under copyright law. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 19 | IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 20 | OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 21 | ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 22 | OTHER DEALINGS IN THE SOFTWARE. 23 | 24 | For more information, please refer to 25 | -------------------------------------------------------------------------------- /experiment_plots/.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore everything in this directory 2 | * 3 | # Except this file 4 | !.gitignore 5 | -------------------------------------------------------------------------------- /kalman.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from utils import gaussian 4 | 5 | 6 | class Gaussian(object): 7 | def __init__(self, mu, sigma): 8 | self.mu = float(mu) 9 | self.sigma = float(sigma) 10 | 11 | def vectorize(self, xlims, step=0.1): 12 | x_array = np.arange(*xlims, step) 13 | y_array = [100 * gaussian(x, self.mu, self.sigma) for x in x_array] 14 | return x_array, y_array 15 | 16 | def bayesian_update(self, other): 17 | sigma_sum = self.sigma**2 + other.sigma**2 18 | self.mu = ((self.sigma**2) * other.mu + (other.sigma**2) * self.mu) / sigma_sum 19 | self.sigma = np.sqrt(((self.sigma * other.sigma)**2) / sigma_sum) 20 | 21 | def __float__(self): 22 | return float(self.mu) 23 | 24 | 25 | class Reading(Gaussian): 26 | 27 | def __init__(self, sensor, truth, timestamp=None): 28 | self.sensor = sensor 29 | self.timestamp = timestamp 30 | self.truth = truth 31 | self.value = sensor.read(truth) 32 | self.color = sensor.color 33 | 34 | self.mu = sensor.predictor(self.value) 35 | self.sigma = sensor.predictor_sigma 36 | 37 | 38 | class Estimate(Gaussian): 39 | 40 | def __init__(self, color='purple'): 41 | self.reading_vector = [] 42 | self.mu = None 43 | self.sigma = None 44 | self.color = color 45 | 46 | def add_reading(self, reading): 47 | self.reading_vector.append(reading) 48 | self.update(reading) 49 | 50 | def reorder(self): 51 | self.reading_vector.sort(key=lambda x: x.timestamp) 52 | 53 | def update(self, reading): 54 | if self.mu is None or self.sigma is None: 55 | self.mu = reading.mu 56 | self.sigma = reading.sigma 57 | else: 58 | self.bayesian_update(reading) 59 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import datetime 3 | import math 4 | import random 5 | from collections import defaultdict 6 | 7 | import matplotlib as mpl 8 | import matplotlib.animation 9 | from matplotlib.ticker import FormatStrFormatter 10 | 11 | import numpy as np 12 | import scipy 13 | import seaborn as sns 14 | import traces 15 | 16 | # local imports 17 | from kalman import Estimate, Reading 18 | from matplotlib import pyplot as plt 19 | from sensor import Sensor 20 | from traincar import TrainCar 21 | 22 | 23 | class SensorAnimation(matplotlib.animation.FuncAnimation): 24 | 25 | def __init__(self, time_array, truth, reading_array, estimate_array): 26 | 27 | self.fig, (self.ax2, self.ax1) = plt.subplots( 28 | 1, 2, sharey=True, 29 | gridspec_kw={"width_ratios":[3, 1]}, 30 | figsize=(8, 4) 31 | ) 32 | plt.tight_layout(pad=2.0) 33 | 34 | self.time_array = time_array 35 | self.estimate_array = estimate_array 36 | 37 | self.ax1.set_ylim(0, 120) 38 | self.ax1.set_xlim(0, 20) 39 | self.ax1.set_xlabel("Probability") 40 | self.ax1.xaxis.set_major_formatter(FormatStrFormatter('%d%%')) 41 | 42 | self.estimate_line = self.ax1.plot( 43 | [], [], color='purple', label='estimate') 44 | self.lines = [] 45 | for sensor in reading_array: 46 | self.lines += self.ax1.plot( 47 | [], [], color=sensor.color, label=sensor.name) 48 | 49 | self.truth_line = self.ax1.hlines(truth[0], 0, 20, color='red', label='Occupancy') 50 | self.ax1.legend() 51 | 52 | self.ax2.plot(time_array, truth, color='red', label='Occupancy') 53 | # self.ax2.set_ylim(0, 150) 54 | self.ax2.set_title("Train car occupancy over time") 55 | self.ax2.set_xlabel("Time (minutes)") 56 | self.ax2.set_ylabel("Occupants") 57 | self.estimate_ts = self.ax2.plot( 58 | [], [], color='purple', label='estimate') 59 | self.fill_lines = self.ax2.fill_between( 60 | [], [], color='purple', alpha=0.5) 61 | 62 | self.truth = truth 63 | self.reading_array = reading_array 64 | 65 | super().__init__( 66 | self.fig, self.update, 67 | frames=len(time_array), 68 | blit=True 69 | ) 70 | 71 | def update(self, i): 72 | """updates frame i of the animation""" 73 | self.ax1.set_title("{}".format( 74 | datetime.timedelta(minutes=self.time_array[i])) 75 | ) 76 | for sensor, line in zip(self.reading_array.keys(), self.lines): 77 | reading = self.reading_array.get(sensor)[i] 78 | x, y = reading.vectorize(self.ax1.get_ylim()) 79 | line.set_data(y, x) 80 | estimate = self.estimate_array[i] 81 | self.estimate_line[0].set_data( 82 | estimate.vectorize(self.ax1.get_ylim())[1], 83 | estimate.vectorize(self.ax1.get_ylim())[0], 84 | ) 85 | 86 | self.truth_line.remove() 87 | self.truth_line = self.ax1.hlines(truth[i], 0, 20, color='red', label='Occupancy') 88 | 89 | self.estimate_ts[0].set_data( 90 | self.time_array[:i], self.estimate_array[:i]) 91 | self.fill_lines.remove() 92 | self.fill_lines = self.ax2.fill_between( 93 | self.time_array[:i], 94 | [e.mu - 2 * e.sigma for e in self.estimate_array[:i]], 95 | [e.mu + 2 * e.sigma for e in self.estimate_array[:i]], 96 | color='purple', 97 | alpha=0.5 98 | ) 99 | 100 | return tuple(self.lines + self.estimate_line + self.estimate_ts + [self.fill_lines] + [self.truth_line]) 101 | 102 | 103 | if __name__ == "__main__": 104 | 105 | # create some crappy sensors 106 | co2_sensor = Sensor("CO$_2$", intersect=350, slope=15, 107 | sigma=10, round_level=500, proc_sigma=30, units="ppm") 108 | # sigma=500, round_level=500, proc_sigma=0) 109 | temp_sensor = Sensor("Temperature", intersect=0, slope=0.25, 110 | sigma=5, round_level=10, proc_sigma=5, units="$^{\circ}$C") 111 | 112 | # put the sensors on a train car 113 | train_car = TrainCar(sensor_array=[co2_sensor, temp_sensor]) 114 | 115 | # run some experiments to model/calibrate the sensors 116 | train_car.run_experiment(datapoints=250) 117 | train_car.plot_experiment(path="experiment_plots") 118 | 119 | # generate some "real" occupancy data 120 | train_car.generate_occupancy() # defaults to 5 stations and 30 minutes 121 | 122 | time_array = np.arange(0, 30, 1.0 / 10) 123 | reading_array = defaultdict(list) 124 | truth = [] 125 | estimate_array = [] 126 | estimate = Estimate() 127 | for t in time_array: 128 | for reading in train_car.read_sensors(t): 129 | reading_array[reading.sensor].append(reading) 130 | estimate.add_reading(reading) 131 | estimate_array.append(copy.deepcopy(estimate)) 132 | # if the last point was in a station 133 | if truth and train_car.occupants_trace[t] != truth[-1]: 134 | estimate = Estimate() 135 | truth.append(train_car.occupants_trace[t]) 136 | 137 | # plt.clf() 138 | # plt.plot(time_array, reading_array[co2_sensor]) 139 | # plt.savefig("co2.png") 140 | 141 | plt.clf() 142 | 143 | animation = SensorAnimation( 144 | time_array, truth, reading_array, estimate_array 145 | ) 146 | animation.save("30minutes.mp4", fps=10, bitrate=1024) 147 | 148 | plt.clf() 149 | plt.xlabel("Number of people in the train car") 150 | plt.ylabel("Probability") 151 | plt.gca().yaxis.set_major_formatter(FormatStrFormatter('%.1f%%')) 152 | 153 | reading_1 = Reading(co2_sensor, 60) 154 | print("reading_1 = ", (reading_1.value, reading_1.mu)) 155 | plt.plot(*reading_1.vectorize((0,120)), color=co2_sensor.color, label="CO$_2$ sensor") 156 | plt.vlines(reading_1, 0, max(reading_1.vectorize((0,120))[1]), linestyles='dashed') 157 | plt.legend() 158 | plt.savefig("reading_plots/1_co2.svg") 159 | 160 | reading_2 = Reading(co2_sensor, 60) 161 | print("reading_2 = ", (reading_2.value, reading_2.mu)) 162 | plt.plot(*reading_2.vectorize((0,120)), color=co2_sensor.color) 163 | plt.vlines(reading_2, 0, max(reading_2.vectorize((0,120))[1]), linestyles='dashed') 164 | plt.savefig("reading_plots/2_co2.svg") 165 | 166 | estimate = Estimate() 167 | estimate.add_reading(reading_1) 168 | estimate.add_reading(reading_2) 169 | estimate_line = plt.plot(*estimate.vectorize((0,120)), color='purple', label="Estimate") 170 | plt.legend() 171 | plt.savefig("reading_plots/3_co2.svg") 172 | 173 | reading_3 = Reading(temp_sensor, 60) 174 | print("reading_3 = ", (reading_3.value, reading_3.mu)) 175 | plt.plot(*reading_3.vectorize((0,120)), color=temp_sensor.color, label="Temperature sensor") 176 | plt.vlines(reading_3, 0, max(reading_3.vectorize((0,120))[1]), linestyles='dashed') 177 | estimate.add_reading(reading_3) 178 | estimate_line[0].remove() 179 | estimate_line = plt.plot(*estimate.vectorize((0,120)), color='purple', label="Estimate") 180 | plt.legend() 181 | plt.savefig("reading_plots/4_co2.svg") 182 | -------------------------------------------------------------------------------- /matplotlibrc: -------------------------------------------------------------------------------- 1 | backend: svg 2 | -------------------------------------------------------------------------------- /reading_plots/.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore everything in this directory 2 | * 3 | # Except this file 4 | !.gitignore 5 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | matplotlib 3 | seaborn 4 | ipywidgets 5 | ipdb 6 | ipython 7 | traces 8 | -------------------------------------------------------------------------------- /sensor.py: -------------------------------------------------------------------------------- 1 | import random 2 | import os 3 | import copy 4 | 5 | import numpy as np 6 | import seaborn as sns 7 | from matplotlib import pyplot as plt 8 | from scipy.stats import linregress 9 | 10 | 11 | from utils import gaussian 12 | 13 | 14 | def plot_linear_fit(ax, x_array, y_array, fit_function, fit_sigma, color, cmap): 15 | xlim = (min(x_array), max(x_array)) 16 | ylim = (min(y_array), max(y_array)) 17 | ax.set_xlim(*xlim) 18 | ax.set_ylim(*ylim) 19 | x_range = np.linspace(*xlim) 20 | y_range = np.linspace(*ylim) 21 | 22 | ax.scatter(x_array, y_array, lw=0, alpha=0.5, color=color) 23 | fit_line = [fit_function(x) for x in x_range] 24 | ax.plot(x_range, fit_line, color=color) 25 | 26 | xx, yy = np.meshgrid(x_range, y_range) 27 | zz = xx + yy 28 | 29 | for i in range(len(x_range)): 30 | for j in range(len(y_range)): 31 | zz[j, i] = gaussian(yy[j, i], fit_function(xx[j, i]), fit_sigma) 32 | 33 | im = ax.imshow( 34 | zz, origin='lower', interpolation='bilinear', 35 | cmap=cmap, alpha=0.5, aspect='auto', 36 | extent=(xlim[0], xlim[-1], ylim[0], ylim[-1]), 37 | vmin=0.0, vmax=gaussian(0, 0, fit_sigma) 38 | ) 39 | 40 | return ax, im 41 | 42 | 43 | class Sensor(object): 44 | 45 | def __init__(self, name, **kwargs): 46 | self.name = name 47 | for key, value in kwargs.items(): 48 | setattr(self, key, value) 49 | 50 | def read(self, variable): 51 | variable = max(0, random.gauss(variable, self.proc_sigma)) 52 | reading = variable * self.slope + self.intersect 53 | return random.gauss(reading, self.sigma) 54 | 55 | def fit(self, data): 56 | self.experiment_data = copy.deepcopy(data) 57 | n_samples = len(data) 58 | 59 | model_slope, model_intercept = np.polyfit( 60 | [o for o, r in data], [r for o, r in data], 1) 61 | 62 | def model(occupants): 63 | return occupants * model_slope + model_intercept 64 | self.model = model 65 | 66 | def predictor(sensor_reading): 67 | return (sensor_reading-model_intercept)/model_slope 68 | self.predictor = predictor 69 | 70 | error = 0.0 71 | for occupants, reading in data: 72 | error += (predictor(reading) - occupants)**2 73 | sigma = np.sqrt(error / (n_samples - 1)) 74 | 75 | self.predictor_sigma = sigma 76 | 77 | def plot_experiment(self, path=""): 78 | color = self.color 79 | data = self.experiment_data 80 | cmap = sns.light_palette(color, as_cmap=True) 81 | 82 | fig, ax = plt.subplots() 83 | occupants, readings = (np.array(array) for array in zip(*data)) 84 | 85 | # ax_left, im_left = plot_linear_fit( 86 | # ax_left, occupants, readings, self.model, self.model_sigma, color, 87 | # cmap) 88 | 89 | ax, im = plot_linear_fit( 90 | ax, readings, occupants, 91 | self.predictor, self.predictor_sigma, 92 | color, cmap 93 | ) 94 | 95 | ax.set_xlabel("{} sensor readout ({})".format(self.name, self.units)) 96 | ax.set_ylabel("Number of train car occupants") 97 | 98 | # cax, kw = mpl.colorbar.make_axes( 99 | # [ax_left, ax_right], location="bottom" 100 | # ) 101 | 102 | # norm = mpl.colors.Normalize(vmin=0, vmax=1) 103 | # cbar = mpl.colorbar.ColorbarBase( 104 | # ax, cmap=cmap, norm=norm, alpha=0.5) 105 | 106 | cbar = plt.colorbar(im, alpha=0.5, extend='neither', ticks=[ 107 | gaussian(3 * self.predictor_sigma, 0, self.predictor_sigma), 108 | gaussian(2 * self.predictor_sigma, 0, self.predictor_sigma), 109 | gaussian(self.predictor_sigma, 0, self.predictor_sigma), 110 | gaussian(0, 0, self.predictor_sigma), 111 | ]) 112 | # cbar.solids.set_edgecolor("face") 113 | 114 | cbar.set_ticklabels( 115 | ['$3 \sigma$', '$2 \sigma$', '$\sigma$', '{:.2%}'.format( 116 | gaussian(0, 0, self.predictor_sigma))], 117 | update_ticks=True 118 | ) 119 | 120 | fig.savefig(os.path.join(path, self.name+".svg")) 121 | -------------------------------------------------------------------------------- /traincar.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | import traces 4 | import seaborn as sns 5 | 6 | from sensor import Sensor 7 | from kalman import Reading 8 | 9 | class TrainCar(object): 10 | max_occupants = 120 11 | occupant_range = range(0, max_occupants + 1) 12 | 13 | def __init__(self, occupants=0, sensor_array=[]): 14 | self.sigma = 0 15 | self.occupants = occupants 16 | self.sensor_array = sensor_array 17 | 18 | def generate_occupancy(self, start=0, end=30, stations=5): 19 | self.occupants_trace = traces.TimeSeries() 20 | self.occupants_trace[start] = self.max_occupants / 2 21 | self.occupants_trace[end] = 0 22 | 23 | # at each station a certain number of people get on or off 24 | minute = start 25 | while minute < end: 26 | minute += random.randint(5, 10) 27 | current_val = self.occupants_trace[minute] 28 | # new_val = max(0, int(random.gauss(current_val, 40))) 29 | new_val = random.randint(0, self.max_occupants) 30 | self.occupants_trace[minute] = new_val 31 | 32 | return self.occupants_trace 33 | 34 | def read_sensors(self, timestamp): 35 | for sensor in self.sensor_array: 36 | occupants = self.occupants_trace[timestamp] 37 | yield Reading(sensor, occupants, timestamp) 38 | 39 | 40 | def plot_experiment(self, **kwargs): 41 | for i, sensor in enumerate(self.sensor_array): 42 | color = sns.color_palette()[i] 43 | sensor.color = color 44 | sensor.plot_experiment(**kwargs) 45 | 46 | def run_experiment(self, datapoints=1000): 47 | """Generates fake sensor data and trains the sensor""" 48 | 49 | for sensor in self.sensor_array: 50 | sensor_data = [] 51 | for _ in range(datapoints): 52 | occupants = random.randrange(self.max_occupants + 1) 53 | sensor_data.append((occupants, sensor.read(occupants))) 54 | sensor.fit(sensor_data) 55 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def gaussian(x, mu, sig): 5 | return np.exp( 6 | -np.power(x - mu, 2.) / (2 * np.power(sig, 2.)) 7 | ) / (sig * np.sqrt(2. * np.pi)) 8 | --------------------------------------------------------------------------------