diff --git a/aerosandbox/atmosphere/atmosphere.py b/aerosandbox/atmosphere/atmosphere.py index 2277170d..0b5a02f0 100644 --- a/aerosandbox/atmosphere/atmosphere.py +++ b/aerosandbox/atmosphere/atmosphere.py @@ -53,6 +53,74 @@ def __repr__(self) -> str: return f"Atmosphere ({altitude_string}, method: '{self.method}')" + def __getitem__(self, index) -> "Atmosphere": + """ + Indexes one item from each attribute of an Atmosphere instance. + Returns a new Atmosphere instance. + + Args: + index: The index that is being called; e.g.,: + >>> first_atmosphere = atmosphere[0] + + Returns: A new Atmosphere instance, where each attribute is subscripted at the given value, if possible. + + """ + l = len(self) + + def get_item_of_attribute(a): + if hasattr(a, "__len__") and hasattr(a, "__getitem__"): + if len(a) == 1: + return a[0] + elif len(a) == l: + return a[index] + else: + try: + return a[index] + except IndexError as e: + raise IndexError(f"A state variable could not be indexed; it has length {len(a)} while the" + f"parent has length {l}.") + else: + return a + + inputs = { + "altitude" : self.altitude, + "temperature_deviation": self.temperature_deviation, + } + + return self.__class__( + **{ + k: get_item_of_attribute(v) + for k, v in inputs.items() + }, + method=self.method, + ) + + def __len__(self): + length = 1 + for v in [ + self.altitude, + self.temperature_deviation, + ]: + if np.length(v) == 1: + try: + v[0] + length = 1 + except (TypeError, IndexError, KeyError) as e: + pass + elif length == 0 or length == 1: + length = np.length(v) + elif length == np.length(v): + pass + else: + raise ValueError("State variables are appear vectorized, but of different lengths!") + return length + + def __array__(self, dtype="O"): + """ + Allows NumPy array creation without infinite recursion in __len__ and __getitem__. + """ + return np.fromiter([self], dtype=dtype).reshape(()) + ### The two primary state variables, pressure and temperature, go here! def pressure(self):