Incorporating classes, objects, and functions into your code will improve its efficiency, readability, and make it easier to extend to other programs down the road. This notebook attempts to take you through many of the concepts associated with creating a class and have it do something useful and interesting.
By C Hummels (Caltech / NSF Fellow)
In [ ]:
import matplotlib.pyplot as plt
import random
import numpy as np
%matplotlib inline
Problem 1a
What do you think should be the attributes and characteristics that should be tracked for a galaxy object? Think critically and make a list of at least six different quantities that would make sense if this were a research project. Please do this prior to starting (or looking at) the example below. Discuss with your code partner.
Provide your answer to Problem 1a here
Problem 1b
Define a galaxy class, and give it an initializer class method. Have your initializer accept arguments for total_mass
, cold_gas_mass
, and stellar_mass
. Have it also accept keyword arguments for age
with a default set to 0. Maybe also include a SFR
and color
in your attributes, but initialize it to 0 and 'red'
respectively. I'll include docstrings for you to help guide you along the way.
Hint: Don't forget to use self
to indicate an object of this class.
In [ ]:
class Galaxy():
"""
Galaxy class for simply representing a galaxy.
"""
def __init__( #complete
Problem 1c
Try creating a galaxy instance of your galaxy class! Create a galaxy object, named after your favorite galaxy, with appropriate values for your arguments and keyword arguments (or as best as you can approximate). There isn't much you can do with it right now, but at least you can be assured that your existing code works. Try printing your object to see what python thinks of it.
In [ ]:
milky_way = Galaxy( #complete
Problem 1d
So Python probably didn't do a very good internal representation of your galaxy. Remember, classes have a few built-in class methods to be defined. One is the __init__
method, so that we know how to initialize the class. Another very useful class method is the __repr__
, which tells us how to represent a class object when we try to print it out. It's very useful for debugging. Please copy your class definition from above, and add the __repr__
method to return a string somehow identifying a galaxy based on its attributes. Now recreate your favorite galaxy, and print it out again and see if it is a bit more readable than before.
In [ ]:
class Galaxy():
"""
Galaxy class for simply representing a galaxy.
"""
def __init__(#complete
Now let's build on our basic galaxy class with some relevant class methods specific to galaxies that enables them to change with time. Obviously, galaxies have a huge number of traits, and we don't currently understand all the ways in which these variables interrelate. But we can make a toy model here just to get the idea of how they might evolve. This is largely the premise of semi-analytic models that people use to study how galaxies might evolve with time.
Problem 2a
Think about what methods you would add to allow your galaxy to change with time using the galaxy attributes defined above and/or other galaxy attributes. Come up with at least three and then discuss with your partner before continuing.
Provide your answer to Problem 2a here
Problem 2b
We want to enhance our general Galaxy
class to enable galaxies to actually evolve with time. Rather than just copying all of the code from our Galaxy
class above, we can subclass the Galaxy
class to inherit all of its attributes, methods, etc! Create a new class named EvolvingGalaxy
below, which uses the Galaxy
class as its "parent" or "super" class. Remember, because you already defined __init__
and __repr__
methods for this class, they are inherited and you don't need to redefine them, unless you want to.
Now try creating an instance of EvolvingGalaxy
and printing it out, just like you did for your favorite galaxy in problem 1c.
In [ ]:
class EvolvingGalaxy( #complete
"""
Galaxy class for representing a galaxy that can evolve over time.
"""
#complete
Problem 2c
Let's start by making an evolve
method for our Galaxy
class in order to evolve the galaxy forward in time. evolve
should expect an argument for time
in years. And at least initially, let's just make this a simple method that just changes the age
of the galaxy by the appropriate amount. We'll add more on to it in a second.
Because we're going to be following the evolution of this galaxy, we should create a current_state
method, which returns a tuple of the attributes of the EvolvingGalaxy
object at that moment. Now as our EvolvingGalaxy
changes with time, we will be able to track it and save that information somewhere.
In [ ]:
class EvolvingGalaxy( #complete
"""
Galaxy class for representing a galaxy that can evolve over time.
"""
#complete
Now that we can evolve our EvolvingGalaxy
and check on its current_state
over time, let's run it forward in time and see the results!
Make a new function (not a class method) called integrate_time
that accepts an EvolvingGalaxy
object, a timestep (some number of years long), and a number of timesteps over which to evolve the galaxy. Step the EvolvingGalaxy
forward in time using your for
loop, one timestep at a time, and store its current_state
at each timestep. You can store this current_state information however you'd like, but I find it easiest to create a NumPy array of size [5, n_timesteps
], and fill it in using the for
loop.
Similarly, make a new function called plot_galaxy_evolution
that accepts the output from your integrate_time
function and uses Matplotlib's pyplot module to plot up the evolution of the EvolvingGalaxy
's properties versus its age.
Finally, create an instance of EvolvingGalaxy, and plug it into your new integrate_time
and plot_galaxy_evolution
functions to integrate for a gigayear, and plot how the EvolvingGalaxy
quantities change with age. Not much going on--sad!
In [ ]:
def integrate_time( #complete
"""
Integrate the time forward for a galaxy and record its state at each timestep; return as array
"""
#complete
def plot_galaxy_evolution( #complete
"""
Plot the evolution of a galaxy from its input data array
"""
#complete
#complete
Problem 2e
How do galaxies evolve with time? They do in lots of ways, but let's make a few specific rules.
Galaxies tend to accrete gas from the intergalactic medium over time. We'll modify the code from above for your EvolvingGalaxy
class to include a accrete_gas_from_IGM
method, which adds to the cold_gas_mass
and total_mass
of the galaxy with time. Let's give it a random component too.
Galaxies form stars, converting cold_gas_mass
to stellar_mass
according to their star formation rate (SFR
). But there must be sufficient cold_gas_mass
to consume, so make sure there is enough for the SFR
to operate for the alloted time
, otherwise the SFR
should drop to 0.
Finally, we haven't actually calculated a SFR
, so let's just create a method, calculate_star_formation_rate
, to choose a random value every time it is called. We'll also scale it in some way by the total_mass
of the galaxy and use the Milky Way as a guide. The Milky Way forms about 1 Msun/year and it has a total_mass
of 1e12 Msun.
Note: You can use the random.random() function to get a random number between 0 and 1.
Include calls to these methods in our evolve
method, and re-generate a galaxy evolution plot like you did in Problem 2d. Since there is a random component, you can try re-running it a few times and see how the behavior changes.
In [ ]:
class EvolvingGalaxy(Galaxy):
"""
Galaxy class for representing a galaxy that can evolve over time.
"""
def current_state( #complete
"""
Return a tuple of the galaxy's total_mass, cold_gas_mass, stellar_mass, age, and SFR
"""
#complete
def calculate_star_formation_rate( #complete
"""
Calculate the star formation rate by taking a random number between 0 and 1
normalized by the galaxy total mass / 1e12;
Also updates the galaxy's color to blue if SFR > 0.01, otherwise color = red
"""
#complete
def accrete_gas_from_IGM( #complete
"""
Allow the galaxy to accrete cold gas from the IGM at a variable rate normalized to
the galaxy's mass
"""
#complete
def form_stars( #complete
"""
Form stars according to the current star formation rate and time available
If unable cold gas, then shut off star formation
"""
#complete
def evolve( #complete
"""
Evolve this galaxy forward for a period time
"""
#complete
#complete
Problem 2f
Based on your knowledge of how galaxies evolve, do these results make sense? Discuss with your coding partner if so/not, and ways in which you might tweak it to work better.
Provide your answer to Problem 2f here
Problem 3a
Let us create a new subclass of EvolveGalaxy
, so that it inherits most of the traits of our previous work. But we will add some new features including position and velocity values in the x-y plane, so that our galaxy can actually move over time.
Create a new class MovingGalaxy
that is a subclass of EvolveGalaxy
. Make a new __init__
method with all of the arguments from it's parent class EvolveGalaxy
, but add the additional arguments of x_position
, y_position
, x_velocity
, y_velocity
, and id_num
.
Note: all arguments must come before any keyword arguments, so the age
kwarg has to come at the very end.
Also, make a new __repr__
method for this class that just includes it's idnum
, x_position
, and y_position
.
To assure your code works, create an instance of a MovingGalaxy
, and print it out.
In [ ]:
class MovingGalaxy(EvolvingGalaxy):
"""
Galaxy class that can evolve and move in the x,y plane
"""
def __init__(self, total_mass, cold_gas_mass, stellar_mass, #complete
# Replace self with super to activate the superclass's methods
super().__init__(total_mass, cold_gas_mass, stellar_mass)
#complete
def __repr__(self):
#complete
#complete
Problem 3b
In order for our MovingGalaxy
to move, we have to give it a few more methods. Copy your class definition from above, and add three more methods.
Create a move
method which accepts a time
argument and updates the MovingGalaxy
's x_position
, and y_position
by accounting for the motion from its x_velocity
and y_velocity
.
Create a calculate_momentum
method that returns a tuple of (x_momentum, y_momentum) for a galaxy.
Finally, create an evolve
method, which accepts a time
argument, and just executes your move
method before calling the superclass's evolve
method (i.e. EvolvingGalaxy
's).
In [ ]:
class MovingGalaxy(EvolvingGalaxy):
"""
This galaxy can move over time in the x,y plane
"""
def __init__(self, total_mass, cold_gas_mass, stellar_mass, x_position, y_position, x_velocity, y_velocity, idnum, age=0):
# Replace self with super to activate the superclass's methods
super().__init__(total_mass, cold_gas_mass, stellar_mass)
self.x_position = x_position
self.y_position = y_position
self.x_velocity = x_velocity
self.y_velocity = y_velocity
self.idnum = idnum
def __repr__(self):
return "Galaxy %i (x = %.0f; y = %.0f)" % (self.idnum, self.x_position, self.y_position)
def move( #complete
def calculate_momentum( #complete
def evolve( #complete
Problem 3c
OK, so we have a MovingGalaxy
that can move, but we need a Universe in which it can move!
So let us create a new Universe
class that can hold one or more galaxies. We will give it x and y limits which define its borders to be 0 and 100, as well as an internal list to hold the galaxies themselves. Make an evolve
method for the Universe
class that moves its time forward some timestep. How does this affect the galaxies? Are there boundaries to your Universe
? What happens when MovingGalaxy
's run into each other? Don't worry about gravity for now, but you could make them merge when they are close enough to each other, or just let them fly through each other. After all, it's your Universe.
It may also be useful to make a method to add a random MovingGalaxy
, so that you can add several to your Universe
during initialization. Start out only adding a few, as it might get expensive!
Make a command that plots the positions of the galaxies with time using matplotlib.scatter()
, and save the corresponding outputs to a files called frame%04i.png
. The next steps allow you to make a movie of these frames.
In [ ]:
class Universe():
"""
A Universe for galaxies to move around in and evolve
"""
def __init__( #complete
def __repr__( #complete
#complete
Problem 3d
Initialize a Universe
with several galaxies in it. Run your Universe
forward in time, and plot the resulting state of the Universe
over time to files called frame%04i.png
. But be careful to not use too many galaxies or to evolve the Universe
forward, as it may be really expensive for your computer. Start small and build from there!
In [ ]:
#complete
Problem 3e
The command below will output a movie file called movie.mp4
by combining all of your frame%04i.png
files and packaging them together at a framerate of 20 FPS. It requires ffmpeg to be installed, but you can do this easily if you have homebrew installed by running brew install ffmpeg
. There are other means to download ffmpeg
online. Otherwise, you can flip through the frames manually to see how the Universe
changes with time.
If you were able to execute this step OK, execute to the next two fields and watch your movie! Again, since there is a random component to these models, you can try running it a few times to see how the results change.
In [ ]:
%%bash
ffmpeg -r 20 -f image2 -i frame%04d.png -vcodec libx264 -pix_fmt yuv420p -crf 25 -y movie.mp4
In [ ]:
%%HTML
<video width="1000" height="1000" controls>
<source src="movie.mp4" type="video/mp4">
</video>
Using the framework you developed for Problems 1-3, add in gravitational effects to calculate the gravitational acceleration of each galaxy with each other galaxy. This is what is known as an N-body integrator, and it can be surprisingly easy to write one. Do this for your Universe
above and run it forward in time.
In [ ]: