def as_proto(self) -> protos.PhysicalState: """Creates a protos.PhysicalState view into all internal data. Expensive. Consider one of the other accessors, which are faster. For example, if you want to iterate over all elements, use __iter__ by doing: for entity in my_physics_state: print(entity.name)""" constructed_protobuf = protos.PhysicalState() constructed_protobuf.CopyFrom(self._proto_state) for entity_data, entity in zip(self, constructed_protobuf.entities): (entity.x, entity.y, entity.vx, entity.vy, entity.heading, entity.spin, entity.fuel, entity.throttle, entity.landed_on, entity.broken) = (entity_data.x, entity_data.y, entity_data.vx, entity_data.vy, entity_data.heading, entity_data.spin, entity_data.fuel, entity_data.throttle, entity_data.landed_on, entity_data.broken) return constructed_protobuf
def load_savefile(file: Path) -> 'data_structures.PhysicsState': """Loads the physics state represented by the input file. If the input file is an OrbitX-style .json file, simply loads it. If the input file is an OrbitV-style .rnd file, tries to interpret it and also loads the adjacent STARSr file to produce a PhysicsState.""" # We shouldn't import orbitv_file_interface at the top of common.py, since # common.py is imported by lots of modules and shouldn't circularly depend # on anything. from orbitx import orbitv_file_interface physics_state: data_structures.PhysicsState logging.getLogger().info(f'Loading savefile {file.resolve()}') assert isinstance(file, Path) if file.suffix.lower() == '.rnd': physics_state = \ orbitv_file_interface.clone_orbitv_state(file) else: if file.suffix.lower() != '.json': logging.getLogger().warning( f'{file} is not a .json file, trying to load it anyways.') with open(file, 'r') as f: data = f.read() read_state = protos.PhysicalState() google.protobuf.json_format.Parse(data, read_state) physics_state = data_structures.PhysicsState(None, read_state) if physics_state.time_acc == 0: physics_state.time_acc = 1 if physics_state.reference == '': physics_state.reference = DEFAULT_REFERENCE if physics_state.target == '': physics_state.target = DEFAULT_TARGET if physics_state.srb_time == 0: physics_state.srb_time = SRB_FULL return physics_state
def __init__(self, y: Optional[np.ndarray], proto_state: protos.PhysicalState): """Collects data from proto_state and y, when y is not None. There are two kinds of values we care about: 1) values that change during simulation (like position, velocity, etc) 2) values that do not change (like mass, radius, name, etc) If both proto_state and y are given, 1) is taken from y and 2) is taken from proto_state. This is a very quick operation. If y is None, both 1) and 2) are taken from proto_state, and a new y vector is generated. This is a somewhat expensive operation.""" assert isinstance(proto_state, protos.PhysicalState) assert isinstance(y, np.ndarray) or y is None # self._proto_state will have positions, velocities, etc for all # entities. DO NOT USE THESE they will be stale. Use the accessors of # this class instead! self._proto_state = protos.PhysicalState() self._proto_state.CopyFrom(proto_state) self._n = len(proto_state.entities) self._entity_names = \ [entity.name for entity in self._proto_state.entities] self._array_rep: np.ndarray if y is None: # We rely on having an internal array representation we can refer # to, so we have to build up this array representation. y = np.empty( len(proto_state.entities) * len(_PER_ENTITY_MUTABLE_FIELDS) + self.N_SINGULAR_ELEMENTS, dtype=self.DTYPE) for field_name, field_n in _FIELD_ORDERING.items(): for entity_index, entity in enumerate(proto_state.entities): proto_value = getattr(entity, field_name) # Internally translate string names to indices, otherwise # our entire y vector will turn into a string vector oh no. # Note this will convert to floats, not integer indices. if field_name == _LANDED_ON: proto_value = self._name_to_index(proto_value) y[self._n * field_n + entity_index] = proto_value y[-2] = proto_state.srb_time y[-1] = proto_state.time_acc self._array_rep = y else: # Take everything except the SRB time, the last element. self._array_rep = y.astype(self.DTYPE) self._proto_state.srb_time = y[self.SRB_TIME_INDEX] self._proto_state.time_acc = y[self.TIME_ACC_INDEX] assert len(self._array_rep.shape) == 1, \ f'y is not 1D: {self._array_rep.shape}' assert (self._array_rep.size - self.N_SINGULAR_ELEMENTS) % \ len(_PER_ENTITY_MUTABLE_FIELDS) == 0, self._array_rep.size assert (self._array_rep.size - self.N_SINGULAR_ELEMENTS) // \ len(_PER_ENTITY_MUTABLE_FIELDS) == len(proto_state.entities), \ f'{self._array_rep.size} mismatches: {len(proto_state.entities)}' np.mod(self.Heading, 2 * np.pi, out=self.Heading) self._entities_with_atmospheres: Optional[List[int]] = None
class PhysicsStateTestCase(unittest.TestCase): """Tests state.PhysicsState accessors and setters.""" proto_state = protos.PhysicalState( timestamp=5, entities=[ protos.Entity(name='First', mass=100, r=200, x=10, y=20, vx=30, vy=40, heading=7, spin=50, fuel=60, throttle=70), protos.Entity(name='Second', mass=101, r=201, artificial=True, x=11, y=21, vx=31, vy=41, heading=2, spin=51, fuel=61, throttle=71, landed_on='First', broken=True) ], engineering=protos.EngineeringState( components=[protos.EngineeringState.Component()] * _N_COMPONENTS, coolant_loops=[protos.EngineeringState.CoolantLoop()] * _N_COOLANT_LOOPS, radiators=[protos.EngineeringState.Radiator()] * _N_RADIATORS)) def test_landed_on(self): """Test that the special .landed_on field is properly set.""" ps = PhysicsState(None, self.proto_state) self.assertEqual(ps['First'].landed_on, '') self.assertEqual(ps['Second'].landed_on, 'First') def test_y_vector_init(self): """Test that initializing with a y-vector uses y-vector values.""" y0 = np.concatenate(( np.array([ 10, 20, # x 30, 40, # y 50, 60, # vx 0, 0, # vy 0, 0, # heading 70, 80, # spin 90, 100, # fuel 0, 0, # throttle 1, -1, # only First is landed on Second 0, 1, # Second is broken common.SRB_EMPTY, 1 # time_acc ]), np.zeros(EngineeringState.N_ENGINEERING_FIELDS))) ps = PhysicsState(y0, self.proto_state) self.assertTrue(np.array_equal(ps.y0(), y0.astype(ps.y0().dtype))) self.assertEqual(ps['First'].landed_on, 'Second') proto_state = ps.as_proto() proto_state.timestamp = 50 self.assertEqual(proto_state.timestamp, 50) self.assertEqual(proto_state.entities[0].fuel, 90) self.assertTrue(proto_state.entities[1].broken) def test_get_set(self): """Test __getitem__ and __setitem__.""" ps = PhysicsState(None, self.proto_state) entity = ps[0] entity.landed_on = 'Second' ps[0] = entity self.assertEqual(ps[0].landed_on, 'Second') def test_entity_view(self): """Test that setting and getting _EntityView attrs propagate.""" ps = PhysicsState(None, self.proto_state) self.assertEqual(ps[0].name, 'First') entity = ps[0] self.assertTrue(isinstance(entity, _EntityView)) self.assertEqual(entity.x, 10) self.assertEqual(entity.y, 20) self.assertEqual(entity.vx, 30) self.assertEqual(entity.vy, 40) self.assertEqual(entity.spin, 50) self.assertEqual(entity.fuel, 60) self.assertEqual(entity.landed_on, '') self.assertEqual(entity.throttle, 70) ps.y0() self.assertEqual(entity.heading, 7 % (2 * np.pi)) ps[0].landed_on = 'Second' self.assertEqual(entity.landed_on, 'Second') entity.x = 500 self.assertEqual(ps[0].x, 500) entity.pos = np.array([55, 66]) self.assertEqual(ps['First'].x, 55) self.assertEqual(ps['First'].y, 66)
def clone_orbitv_state(rnd_path: Path) -> PhysicsState: """Gathers information from an OrbitV savefile, in the *.RND format, as well as a STARSr file. Returns a PhysicsState representing everything we can read.""" proto_state = protos.PhysicalState() starsr_path = rnd_path.parent / 'STARSr' log.info(f"Reading OrbitV file {rnd_path}, " f"using info from adjacent STARSr file {starsr_path}") hab_index: Optional[int] = None ayse_index: Optional[int] = None with open(starsr_path, 'r') as starsr_file: starsr = csv.reader(starsr_file, delimiter=',') background_star_line = next(starsr) while len(background_star_line) == 3: background_star_line = next(starsr) # After the last while loop, the value of background_star_line is # actually a gravity_pair line (a pair of indices, which we also don't # care about). gravity_pair_line = background_star_line while len(gravity_pair_line) == 2: gravity_pair_line = next(starsr) entity_constants_line = gravity_pair_line while len(entity_constants_line) == 6: proto_entity = proto_state.entities.add() # Cast all the elements on a line to floats. # Some strings will be the form '1.234D+04', convert these too. (colour, mass, radius, atmosphere_thickness, atmosphere_scaling, atmosphere_height) \ = map(_string_to_float, entity_constants_line) mass = max(1, mass) proto_entity.mass = mass proto_entity.r = radius if atmosphere_thickness and atmosphere_scaling: # Only set atmosphere fields if they both have a value. proto_entity.atmosphere_thickness = atmosphere_thickness proto_entity.atmosphere_scaling = atmosphere_scaling entity_constants_line = next(starsr) # We skip a line here. It's the timestamp line, but the .RND file will # have more up-to-date timestamp info. entity_positions_line = next(starsr) # We don't care about entity positions, we'll get this from the .RND # file as well. while len(entity_positions_line) == 6: entity_positions_line = next(starsr) entity_name_line = entity_positions_line entity_index = 0 while len(entity_name_line) == 1: name = entity_name_line[0].strip() proto_state.entities[entity_index].name = name if name == common.HABITAT: hab_index = entity_index elif name == common.AYSE: ayse_index = entity_index entity_index += 1 entity_name_line = next(starsr) assert proto_state.entities[0].name == common.SUN, ( 'We assume that the Sun is the first OrbitV entity, this has ' 'implications for how we populate entity positions.') assert hab_index is not None, \ "We didn't see the Habitat in the STARSr file!" assert ayse_index is not None, \ "We didn't see AYSE in the STARSr file!" hab = proto_state.entities[hab_index] ayse = proto_state.entities[ayse_index] with open(rnd_path, 'rb') as rnd: check_byte = rnd.read(1) rnd.read(len("ORBIT5S ")) craft_throttle = _read_single(rnd) / 100 # I don't know what Vflag and Aflag do. _read_int(rnd), _read_int(rnd) navmode = Navmode(SFLAG_TO_NAVMODE[_read_int(rnd)]) if navmode is not None: proto_state.navmode = navmode.value _read_double(rnd) # This is drag, we calculate this ourselves. _read_double(rnd) # This is zoom, we ignore this as well. hab.heading = _read_single(rnd) # Skip centre, ref, and targ information. _read_int(rnd), _read_int(rnd), _read_int(rnd) # We don't track if trails are set. _read_int(rnd) drag_profile = _read_single(rnd) proto_state.parachute_deployed = (drag_profile > 0.002) current_srb_output = _read_single(rnd) if current_srb_output != 0: proto_state.srb_time = common.SRB_BURNTIME # We don't care about colour trails or how times are displayed. _read_int(rnd), _read_int(rnd) _read_double(rnd), _read_double(rnd) # No time acc info either. _read_single(rnd) # Ignore some RCPS info _read_int(rnd) # Eflag? Something that gassimv.bas sets. year = _read_int(rnd) day = _read_int(rnd) hour = _read_int(rnd) minute = _read_int(rnd) second = _read_double(rnd) timestamp = datetime( year, 1, 1, hour=hour, minute=minute, second=int(second)) + timedelta(day - 1) proto_state.timestamp = timestamp.timestamp() ayse.heading = _read_single(rnd) _read_int(rnd) # AYSEscrape seems to do nothing. # TODO: Once we implement wind, fill in these fields. _read_single(rnd), _read_single(rnd) hab.spin = numpy.radians(_read_single(rnd)) docked_constant = _read_int(rnd) if docked_constant != 0: hab.landed_on = common.AYSE _read_single(rnd) # Rad shields, overwritten by enghabv.bas. # TODO: once module is implemented, have it be launched here. _read_int(rnd) # module_state = MODFLAG_TO_MODSTATE[modflag] _read_single(rnd), _read_single(rnd) # AYSEdist and OCESSdist. hab.broken = bool(_read_int(rnd)) ayse.broken = bool(_read_int(rnd)) # TODO: Once we implement nav malfunctions, set this field. _read_single(rnd) # Max thrust, ignored because we'll read it from ORBITSSE.rnd _read_single(rnd) # TODO: Once we implement nav malfunctions, set this field. # Also, apparently this is the same field as two fields ago? _read_single(rnd) # TODO: Once we implement wind, fill in these fields. _read_long(rnd) # TODO: There is some information required from MST.rnd to implement # this field (LONGtarg). _read_single(rnd) # We calculate pressure and agrav. _read_single(rnd), _read_single(rnd) # Note, we skip the first entity, the Sun, since OrbitV does. for i in range(1, min(39, len(proto_state.entities))): entity = proto_state.entities[i] entity.x, entity.y, entity.vx, entity.vy = (_read_double(rnd), _read_double(rnd), _read_double(rnd), _read_double(rnd)) hab.fuel = _read_single(rnd) ayse.fuel = _read_single(rnd) check_byte_2 = rnd.read(1) if check_byte != check_byte_2: log.warning('RND file had inconsistent check bytes: ' f'{check_byte} != {check_byte_2}') # Delete any entity with a (0, 0) position that isn't the Sun. # TODO: We also delete the OCESS launch tower, but once we implement # OCESS we should no longer do this. entity_index = 0 while entity_index < len(proto_state.entities): proto_entity = proto_state.entities[entity_index] if round(proto_entity.x) == 0 and round(proto_entity.y) == 0 and \ proto_entity.name != common.SUN or \ proto_entity.name == common.OCESS: del proto_state.entities[entity_index] else: entity_index += 1 hab.artificial = True ayse.artificial = True # Habitat and AYSE masses are hardcoded in OrbitV. hab.mass = 275000 ayse.mass = 20000000 orbitx_state = PhysicsState(None, proto_state) orbitx_state[common.HABITAT] = Entity(hab) orbitx_state[common.AYSE] = Entity(ayse) craft = orbitx_state.craft_entity() craft.throttle = craft_throttle log.debug(f'Interpreted {rnd_path} and {starsr_path} into this state:') log.debug(repr(orbitx_state)) orbitx_state = _separate_landed_entities(orbitx_state) return orbitx_state
def __init__(self, y: Optional[np.ndarray], proto_state: protos.PhysicalState): """Collects data from proto_state and y, when y is not None. There are two kinds of values we care about: 1) values that change during simulation (like position, velocity, etc) 2) values that do not change (like mass, radius, name, etc) If both proto_state and y are given, 1) is taken from y and 2) is taken from proto_state. This is a very quick operation. If y is None, both 1) and 2) are taken from proto_state, and a new y vector is generated. This is a somewhat expensive operation.""" assert isinstance(proto_state, protos.PhysicalState) assert isinstance(y, np.ndarray) or y is None # self._proto_state will have positions, velocities, etc for all # entities. DO NOT USE THESE they will be stale. Use the accessors of # this class instead! self._proto_state = protos.PhysicalState() self._proto_state.CopyFrom(proto_state) self._n = len(proto_state.entities) self._entity_names = \ [entity.name for entity in self._proto_state.entities] self._array_rep: np.ndarray if y is None: # We rely on having an internal array representation we can refer # to, so we have to build up this array representation. self._array_rep = np.empty( len(proto_state.entities) * len(_PER_ENTITY_MUTABLE_FIELDS) + self.N_SINGULAR_ELEMENTS + EngineeringState.N_ENGINEERING_FIELDS, dtype=self.DTYPE) for field_name, field_n in _ENTITY_FIELD_ORDER.items(): for entity_index, entity in enumerate(proto_state.entities): proto_value = getattr(entity, field_name) # Internally translate string names to indices, otherwise # our entire y vector will turn into a string vector oh no. # Note this will convert to floats, not integer indices. if field_name == _LANDED_ON: proto_value = self._name_to_index(proto_value) self._array_rep[self._n * field_n + entity_index] = proto_value self._array_rep[self.SRB_TIME_INDEX] = proto_state.srb_time self._array_rep[self.TIME_ACC_INDEX] = proto_state.time_acc # It's IMPORTANT that we pass in self._array_rep, because otherwise the numpy # array will be copied and EngineeringState won't be modifying our numpy array. self.engineering = EngineeringState( self._array_rep[self.ENGINEERING_START_INDEX:], self._proto_state.engineering, parent_state=self, populate_array=True) else: self._array_rep = y.astype(self.DTYPE) self._proto_state.srb_time = y[self.SRB_TIME_INDEX] self._proto_state.time_acc = y[self.TIME_ACC_INDEX] self.engineering = EngineeringState( self._array_rep[self.ENGINEERING_START_INDEX:], self._proto_state.engineering, parent_state=self, populate_array=False) assert len(self._array_rep.shape) == 1, \ f'y is not 1D: {self._array_rep.shape}' n_entities = len(proto_state.entities) assert self._array_rep.size == ( n_entities * len(_PER_ENTITY_MUTABLE_FIELDS) + self.N_SINGULAR_ELEMENTS + EngineeringState.N_ENGINEERING_FIELDS) np.mod(self.Heading, 2 * np.pi, out=self.Heading) self._entities_with_atmospheres: Optional[List[int]] = None