Skip to content

Commit c295651

Browse files
committed
MNT: adapt discretization schema to RocketPy encoder.
1 parent 8450082 commit c295651

File tree

9 files changed

+33
-112
lines changed

9 files changed

+33
-112
lines changed

src/services/environment.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from rocketpy.environment.environment import Environment as RocketPyEnvironment
66
from src.models.environment import EnvironmentModel
77
from src.views.environment import EnvironmentSimulation
8-
from src.utils import rocketpy_encoder, DiscretizeConfig
8+
from src.utils import rocketpy_encoder
99

1010

1111
class EnvironmentService:
@@ -50,9 +50,7 @@ def get_environment_simulation(self) -> EnvironmentSimulation:
5050
EnvironmentSimulation
5151
"""
5252

53-
attributes = rocketpy_encoder(
54-
self.environment, DiscretizeConfig.for_environment()
55-
)
53+
attributes = rocketpy_encoder(self.environment)
5654
env_simulation = EnvironmentSimulation(**attributes)
5755
return env_simulation
5856

src/services/flight.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from src.services.rocket import RocketService
99
from src.models.flight import FlightModel
1010
from src.views.flight import FlightSimulation
11-
from src.utils import rocketpy_encoder, DiscretizeConfig
11+
from src.utils import rocketpy_encoder
1212

1313

1414
class FlightService:
@@ -55,9 +55,7 @@ def get_flight_simulation(self) -> FlightSimulation:
5555
Returns:
5656
FlightSimulation
5757
"""
58-
attributes = rocketpy_encoder(
59-
self.flight, DiscretizeConfig.for_flight()
60-
)
58+
attributes = rocketpy_encoder(self.flight)
6159
flight_simulation = FlightSimulation(**attributes)
6260
return flight_simulation
6361

src/services/motor.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
from src.models.sub.tanks import TankKinds
1818
from src.models.motor import MotorKinds, MotorModel
1919
from src.views.motor import MotorSimulation
20-
from src.utils import rocketpy_encoder, DiscretizeConfig
20+
from src.utils import rocketpy_encoder
2121

2222

2323
class MotorService:
@@ -140,7 +140,7 @@ def get_motor_simulation(self) -> MotorSimulation:
140140
Returns:
141141
MotorSimulation
142142
"""
143-
attributes = rocketpy_encoder(self.motor, DiscretizeConfig.for_motor())
143+
attributes = rocketpy_encoder(self.motor)
144144
motor_simulation = MotorSimulation(**attributes)
145145
return motor_simulation
146146

src/services/rocket.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
from src.models.sub.aerosurfaces import NoseCone, Tail, Fins
1818
from src.services.motor import MotorService
1919
from src.views.rocket import RocketSimulation
20-
from src.utils import rocketpy_encoder, DiscretizeConfig
20+
from src.utils import rocketpy_encoder
2121

2222

2323
class RocketService:
@@ -107,9 +107,7 @@ def get_rocket_simulation(self) -> RocketSimulation:
107107
Returns:
108108
RocketSimulation
109109
"""
110-
attributes = rocketpy_encoder(
111-
self.rocket, DiscretizeConfig.for_rocket()
112-
)
110+
attributes = rocketpy_encoder(self.rocket)
113111
rocket_simulation = RocketSimulation(**attributes)
114112
return rocket_simulation
115113

src/utils.py

Lines changed: 13 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -2,112 +2,39 @@
22
import io
33
import logging
44
import json
5-
import copy
65

7-
from typing import NoReturn, Tuple
6+
from typing import NoReturn
87

9-
from rocketpy import Function
108
from rocketpy._encoders import RocketPyEncoder
119
from starlette.datastructures import Headers, MutableHeaders
1210
from starlette.types import ASGIApp, Message, Receive, Scope, Send
1311

1412
logger = logging.getLogger(__name__)
1513

1614

17-
class DiscretizeConfig:
18-
"""
19-
Configuration class for RocketPy function discretization.
20-
21-
This class allows easy configuration of discretization parameters
22-
for different types of RocketPy objects and their callable attributes.
23-
"""
24-
25-
def __init__(
26-
self, bounds: Tuple[float, float] = (0, 10), samples: int = 200
27-
):
28-
self.bounds = bounds
29-
self.samples = samples
30-
31-
@classmethod
32-
def for_environment(cls) -> 'DiscretizeConfig':
33-
return cls(bounds=(0, 50000), samples=100)
34-
35-
@classmethod
36-
def for_motor(cls) -> 'DiscretizeConfig':
37-
return cls(bounds=(0, 10), samples=150)
38-
39-
@classmethod
40-
def for_rocket(cls) -> 'DiscretizeConfig':
41-
return cls(bounds=(0, 1), samples=100)
42-
43-
@classmethod
44-
def for_flight(cls) -> 'DiscretizeConfig':
45-
return cls(bounds=(0, 30), samples=200)
46-
47-
48-
def rocketpy_encoder(obj, config: DiscretizeConfig = DiscretizeConfig()):
15+
def rocketpy_encoder(obj):
4916
"""
5017
Encode a RocketPy object using official RocketPy encoders.
5118
52-
This function creates a copy of the object, discretizes callable Function
53-
attributes on the copy, and then uses RocketPy's official RocketPyEncoder for
54-
complete object serialization. The original object remains unchanged.
19+
This function uses RocketPy's official RocketPyEncoder for complete
20+
object serialization.
5521
5622
Args:
5723
obj: RocketPy object (Environment, Motor, Rocket, Flight)
58-
config: DiscretizeConfig object with discretization parameters (optional)
5924
6025
Returns:
6126
Dictionary of encoded attributes
6227
"""
6328

64-
# Create a copy to avoid mutating the original object
65-
obj_copy = copy.deepcopy(obj)
66-
67-
for attr_name in dir(obj_copy):
68-
if attr_name.startswith('_'):
69-
continue
70-
71-
try:
72-
attr_value = getattr(obj_copy, attr_name)
73-
except Exception:
74-
continue
75-
76-
if callable(attr_value) and isinstance(attr_value, Function):
77-
try:
78-
discretized_func = Function(attr_value.source)
79-
discretized_func.set_discrete(
80-
lower=config.bounds[0],
81-
upper=config.bounds[1],
82-
samples=config.samples,
83-
mutate_self=True,
84-
)
85-
86-
setattr(obj_copy, attr_name, discretized_func)
87-
88-
except Exception as e:
89-
logger.warning(f"Failed to discretize {attr_name}: {e}")
90-
91-
try:
92-
json_str = json.dumps(
93-
obj_copy,
94-
cls=RocketPyEncoder,
95-
include_outputs=True,
96-
include_function_data=True,
97-
)
98-
return json.loads(json_str)
99-
except Exception as e:
100-
logger.warning(f"Failed to encode with RocketPyEncoder: {e}")
101-
attributes = {}
102-
for attr_name in dir(obj_copy):
103-
if not attr_name.startswith('_'):
104-
try:
105-
attr_value = getattr(obj_copy, attr_name)
106-
if not callable(attr_value):
107-
attributes[attr_name] = str(attr_value)
108-
except Exception:
109-
continue
110-
return attributes
29+
json_str = json.dumps(
30+
obj,
31+
cls=RocketPyEncoder,
32+
include_outputs=True,
33+
include_function_data=True,
34+
discretize=True,
35+
pickle_callables=False,
36+
)
37+
return json.loads(json_str)
11138

11239

11340
class RocketPyGZipMiddleware:

src/views/environment.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ class EnvironmentSimulation(ApiBaseView):
1414
any new attributes that might be encoded.
1515
"""
1616

17-
model_config = ConfigDict(extra='allow', arbitrary_types_allowed=True)
17+
model_config = ConfigDict(extra='ignore', arbitrary_types_allowed=True)
1818

1919
message: str = "Environment successfully simulated"
2020

@@ -35,9 +35,9 @@ class EnvironmentSimulation(ApiBaseView):
3535
initial_hemisphere: Optional[str] = None
3636
initial_ew: Optional[str] = None
3737
max_expected_height: Optional[int] = None
38-
date: Optional[datetime] = datetime.today() + timedelta(days=1)
39-
local_date: Optional[datetime] = datetime.today() + timedelta(days=1)
40-
datetime_date: Optional[datetime] = datetime.today() + timedelta(days=1)
38+
date: Optional[Any] = datetime.today() + timedelta(days=1)
39+
local_date: Optional[Any] = datetime.today() + timedelta(days=1)
40+
datetime_date: Optional[Any] = datetime.today() + timedelta(days=1)
4141

4242
# Function attributes (discretized by rocketpy_encoder, serialized by RocketPyEncoder)
4343
ellipsoid: Optional[Any] = None

src/views/flight.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from src.views.environment import EnvironmentSimulation
77

88

9-
class FlightSimulation(RocketSimulation, EnvironmentSimulation):
9+
class FlightSimulation(ApiBaseView):
1010
"""
1111
Flight simulation view that handles dynamically encoded RocketPy Flight attributes.
1212
@@ -16,7 +16,7 @@ class FlightSimulation(RocketSimulation, EnvironmentSimulation):
1616
any new attributes that might be encoded.
1717
"""
1818

19-
model_config = ConfigDict(extra='allow', arbitrary_types_allowed=True)
19+
model_config = ConfigDict(extra='ignore', arbitrary_types_allowed=True)
2020

2121
message: str = "Flight successfully simulated"
2222

src/views/motor.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ class MotorSimulation(ApiBaseView):
1313
any new attributes that might be encoded.
1414
"""
1515

16-
model_config = ConfigDict(extra='allow', arbitrary_types_allowed=True)
16+
model_config = ConfigDict(extra='ignore', arbitrary_types_allowed=True)
1717

1818
message: str = "Motor successfully simulated"
1919

src/views/rocket.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
1-
from typing import Optional, Any
1+
from typing import Optional, Any, List, Tuple
22
from pydantic import ConfigDict
33
from src.models.rocket import RocketModel
44
from src.views.interface import ApiBaseView
55
from src.views.motor import MotorView, MotorSimulation
66

77

8-
class RocketSimulation(MotorSimulation):
8+
class RocketSimulation(ApiBaseView):
99
"""
1010
Rocket simulation view that handles dynamically encoded RocketPy Rocket attributes.
1111
@@ -15,16 +15,16 @@ class RocketSimulation(MotorSimulation):
1515
attributes that might be encoded.
1616
"""
1717

18-
model_config = ConfigDict(extra='allow', arbitrary_types_allowed=True)
18+
model_config = ConfigDict(extra='ignore', arbitrary_types_allowed=True)
1919

2020
message: str = "Rocket successfully simulated"
2121

2222
# Core Rocket attributes (always present)
2323
radius: Optional[float] = None
2424
mass: Optional[float] = None
2525
inertia: Optional[tuple] = None
26-
power_off_drag: Optional[float] = None
27-
power_on_drag: Optional[float] = None
26+
power_off_drag: Optional[List[Tuple[float, float]]] = None
27+
power_on_drag: Optional[List[Tuple[float, float]]] = None
2828
center_of_mass_without_motor: Optional[float] = None
2929
coordinate_system_orientation: Optional[str] = None
3030
parachutes: Optional[list] = None

0 commit comments

Comments
 (0)