doc/interactive_tutorial/05-Attributes_and_functions_on_particle_collections.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from matplotlib import pyplot\n",
    "import numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AMUSE pre-defines a number of calculcated attributes on particle sets, such as the kinetic energy of the particles in the set. These calculated attributes are used often and provide a sufficient set to start out with, but they do not define a *complete* set. It's possible to define your own attributes and extend the attributes on a particle set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from amuse.lab import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown in the previous example, you can create a particle set by specifying the number of particles and setting their attributes. You can also create a particle set by using an inital condition function. For stellar clusters the commonly used plummer and king models are available. For this tutorial we will start with a king model. Global clusters created with a king model need the number of stars in the cluster and a dimensionless depth parameter that determines the depth of the potential well in the center of the cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particles = new_king_model(1000, 3)\n",
    "print(particles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Common properties for a stellar cluster are its  center of mass position, total kinetic energy and potential energy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"center of mass\", particles.center_of_mass())\n",
    "print(\"kinetic energy\", particles.kinetic_energy())\n",
    "print(\"potential energy\", particles.potential_energy(G=nbody_system.G))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the potential energy calculation we need to specify the gravitational constant, as the default value will use the gavitational constant in S.I. units and we are working in nbody units for this tutorial.\n",
    "\n",
    "In N-body calculations and reporting, the kinetic and potential energy of a set of stars is often scaled to exactly 0.25 and -0.5 respectively. AMUSE also has a function for this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particles.scale_to_standard()\n",
    "print(\"kinetic energy\", particles.kinetic_energy())\n",
    "print(\"potential energy\", particles.potential_energy(G=nbody_system.G))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note that the potential energy and scaling calculations are implemented as order N-squared operations*\n",
    "\n",
    "Attributes of particle sets are always 1 dimensional by default, an array with a single value per particle attribute. But for some attributes it is easier to work with a 2d set, an array with multiple values (or an array of values) per particle attribute. For example, the positions of all particles. These attributes are called vector-attributes and are defined as a combination of 2 or more simple attributes. \n",
    "\n",
    "The position attribute combines the values of the `x`, `y` and `z` attributes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(particles[0].x)\n",
    "print(particles[0].y)\n",
    "print(particles[0].z)\n",
    "print(particles[0].position)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Other common vector attributes are `velocity` (combination of `vx`,`vy`,`vz`) and `acceleration` (combination of `ax`,`ay`,`az`).\n",
    "\n",
    "You can set the value of a position attribute and the underlying x, y or z attributes will be changed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particles[0].position = [0, 0.1, 0.2] | nbody_system.length\n",
    "print(particles[0].x)\n",
    "print(particles[0].y)\n",
    "print(particles[0].z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can set the value of the x, y or z attribute and the position will change (as the position is just a combination of these attributes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particles[0].x = 0.3 | nbody_system.length\n",
    "print(particles[0].position)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You cannot change an item in the position array and thereby change the x, y, or z positions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this will not change anything in the particles set as the position is a copy\n",
    "particles[0].position[0] = 0.5 | nbody_system.length\n",
    "\n",
    "print(particles[0].x)\n",
    "print(particles[0].position)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use the position attribute on the entire set. Let's print the positions of the first 10 particles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(particles.position[0:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also use the position attribute to set values for the entire set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the position of all particles in the set to the same value\n",
    "particles.position = [0.1, 0.2, 0.3] | nbody_system.length\n",
    "\n",
    "print(particles.position[0:10])\n",
    "print(particles.x[0:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining a new vector attribute is done by calling the `add_vector_attribute` or `add_global_vector_attribute`. The first call will define the attribute on the particle set and not on any other set. The second call will define the attribute on the particle set and any future sets created in the script. (The second call is used in the amuse framework itself to define the `position`, `velocity` and `acceleration` attributes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particles.add_vector_attribute('position2d', ('x', 'y'))\n",
    "print(particles[0].position2d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you enter `particles.add_` and press tab you'll notice two other function besides the `add_vector_attribute` function; `add_calculated_attribute` will create an attribute where the values are calculated based on other attributes, `add_function_attribute` will create a function on the set that gets the set and optional function parameters. These function also have global versions (`add_global_...`). The `add_global_function_attribute` call is used in the AMUSE framework to implement the `kinetic_energy` and `potential_energy` functions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particles.add_function_attribute(\n",
    "    'calculate_mean_mass',\n",
    "    lambda particles: particles.mass.sum() / len(particles)\n",
    ")\n",
    "print(particles.calculate_mean_mass())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}