from __future__ import annotations import math from dataclasses import dataclass from .models import ReducerInstance, StageDerived CASCADE_INTERSTAGE_COMPONENT_IDS = { "cascade_input_shaft", "cascade_output_shaft", "interstage_shaft", "interstage_coupler", "interstage_bearing_1", "interstage_bearing_2", "interstage_support_1", "interstage_support_2", } def stage_carrier_plate_thickness_mm(stage_gap_mm: float) -> float: return max(0.8, min(1.2, max(0.8, stage_gap_mm * 0.45))) def sun_bore_radius_from_derived(derived: StageDerived) -> float: return max(0.8, min(2.2, derived.sun_pitch_radius_mm * 0.35)) @dataclass(frozen=True) class CascadeInterstageDesign: source_component_id: str target_component_id: str shaft_diameter_mm: float shaft_bottom_z_mm: float shaft_top_z_mm: float s1_output_axis_z_mm: float s2_input_axis_z_mm: float coupler_inner_diameter_mm: float coupler_outer_diameter_mm: float coupler_width_mm: float coupler_center_z_mm: float bearing_inner_diameter_mm: float bearing_outer_diameter_mm: float bearing_width_mm: float bearing_ball_diameter_mm: float bearing_ball_count: int bearing_centers_z_mm: dict[str, float] support_inner_diameter_mm: float support_outer_diameter_mm: float support_width_mm: float support_centers_z_mm: dict[str, float] housing_inner_diameter_mm: float housing_outer_diameter_mm: float s1_carrier_top_z_mm: float s2_sun_bottom_z_mm: float s2_carrier_top_z_mm: float s2_sun_bore_diameter_mm: float @property def shaft_span_mm(self) -> float: return self.shaft_top_z_mm - self.shaft_bottom_z_mm def _stage_by_id(instance: ReducerInstance) -> dict[str, StageDerived]: return {stage.stage_id: stage for stage in (instance.derived.stages or [])} def compute_cascade_interstage_design(instance: ReducerInstance) -> CascadeInterstageDesign: if instance.topology_family != "simple_2k_h_cascade": raise ValueError("cascade interstage design requires simple_2k_h_cascade") if not instance.parameters.stages or len(instance.parameters.stages) != 2: raise ValueError("cascade interstage design requires exactly two stage parameter sets") if not instance.requirement.stages or len(instance.requirement.stages) != 2: raise ValueError("cascade interstage design requires exactly two stage requirements") stage_derived = _stage_by_id(instance) s1 = instance.parameters.stages[0] s2 = instance.parameters.stages[1] s1d = stage_derived[s1.stage_id] s2d = stage_derived[s2.stage_id] stage_gap = float((instance.derived.axial_stack or {}).get("stage_gap_mm", 2.0)) s1_plate_thickness = stage_carrier_plate_thickness_mm(stage_gap) s2_plate_thickness = stage_carrier_plate_thickness_mm(stage_gap) s1_carrier_top = s1d.z_offset_mm + s1.face_width_mm + 0.25 + s1_plate_thickness s2_sun_bottom = s2d.z_offset_mm available_gap = max(0.35, s2_sun_bottom - s1_carrier_top) s2_sun_bore_radius = sun_bore_radius_from_derived(s2d) shaft_radius = min( max(0.45, s2_sun_bore_radius * 0.72), max(0.35, s2_sun_bore_radius - 0.20), ) bearing_inner_radius = shaft_radius + max(0.08, s2.module_mm * 0.18) bearing_outer_radius = bearing_inner_radius + max(0.35, s2.module_mm * 0.90) bearing_width = max(0.28, min(0.60, available_gap * 0.45)) bearing_ball_diameter = max( 0.12, min( (bearing_outer_radius - bearing_inner_radius) * 0.62, bearing_width * 0.42, ), ) ball_pitch_radius = (bearing_inner_radius + bearing_outer_radius) / 2.0 estimated_ball_count = int((2.0 * math.pi * ball_pitch_radius) / (bearing_ball_diameter * 1.5)) ball_count = max(8, min(18, estimated_ball_count)) bearing_1_center = s1_carrier_top + available_gap * 0.35 s2_carrier_top = s2d.z_offset_mm + s2.face_width_mm + 0.25 + s2_plate_thickness bearing_2_center = s2_carrier_top + 0.35 + bearing_width / 2.0 coupler_width = max(0.20, min(0.40, available_gap * 0.28)) coupler_center = s2_sun_bottom - coupler_width / 2.0 - 0.05 support_inner_radius = bearing_outer_radius + max(0.08, s2.module_mm * 0.20) max_ring_outer_radius = max(stage.ring_outer_diameter_mm / 2.0 for stage in stage_derived.values()) housing_inner_radius = max_ring_outer_radius + 2.5 housing_outer_radius = max_ring_outer_radius + 5.0 support_outer_radius = housing_inner_radius - 0.25 coupler_inner_radius = shaft_radius + max(0.04, s2.module_mm * 0.10) coupler_outer_radius = min( support_inner_radius - max(0.08, s2.module_mm * 0.16), s2_sun_bore_radius + max(0.25, s2.module_mm * 0.75), ) coupler_outer_radius = max(coupler_outer_radius, coupler_inner_radius + max(0.20, s2.module_mm * 0.40)) shaft_bottom = max(s1d.z_offset_mm, s1.face_width_mm - 0.50) shaft_top = bearing_2_center + bearing_width / 2.0 + 0.40 first_requirement, second_requirement = instance.requirement.stages source_component_id = f"{first_requirement.stage_id}_{first_requirement.output}" target_component_id = f"{second_requirement.stage_id}_{second_requirement.input}" return CascadeInterstageDesign( source_component_id=source_component_id, target_component_id=target_component_id, shaft_diameter_mm=2.0 * shaft_radius, shaft_bottom_z_mm=shaft_bottom, shaft_top_z_mm=shaft_top, s1_output_axis_z_mm=s2d.z_offset_mm + s2.face_width_mm / 2.0, s2_input_axis_z_mm=s2d.z_offset_mm + s2.face_width_mm / 2.0, coupler_inner_diameter_mm=2.0 * coupler_inner_radius, coupler_outer_diameter_mm=2.0 * coupler_outer_radius, coupler_width_mm=coupler_width, coupler_center_z_mm=coupler_center, bearing_inner_diameter_mm=2.0 * bearing_inner_radius, bearing_outer_diameter_mm=2.0 * bearing_outer_radius, bearing_width_mm=bearing_width, bearing_ball_diameter_mm=bearing_ball_diameter, bearing_ball_count=ball_count, bearing_centers_z_mm={ "interstage_bearing_1": bearing_1_center, "interstage_bearing_2": bearing_2_center, }, support_inner_diameter_mm=2.0 * support_inner_radius, support_outer_diameter_mm=2.0 * support_outer_radius, support_width_mm=bearing_width, support_centers_z_mm={ "interstage_support_1": bearing_1_center, "interstage_support_2": bearing_2_center, }, housing_inner_diameter_mm=2.0 * housing_inner_radius, housing_outer_diameter_mm=2.0 * housing_outer_radius, s1_carrier_top_z_mm=s1_carrier_top, s2_sun_bottom_z_mm=s2_sun_bottom, s2_carrier_top_z_mm=s2_carrier_top, s2_sun_bore_diameter_mm=2.0 * s2_sun_bore_radius, ) def cascade_interstage_component_ids(instance: ReducerInstance) -> set[str]: if instance.topology_family != "simple_2k_h_cascade": return set() return set(CASCADE_INTERSTAGE_COMPONENT_IDS)