mohr_3d.py

#!/usr/bin/env python3
"""
Mohr's Circle in 3D - Python SVG Generator
Generates Mohr's circle visualization from 3D stress tensor
"""

import math
import sys


class MohrCircle3D:
    """Calculate and visualize Mohr's circle in 3D"""

    def __init__(self, s_x, s_y, s_z, t_xy, t_xz, t_yz, unit="MPa"):
        """
        Initialize with stress components

        Args:
            s_x, s_y, s_z: Normal stresses
            t_xy, t_xz, t_yz: Shear stresses
            unit: Unit label (default: "MPa")
        """
        self.s_x = s_x
        self.s_y = s_y
        self.s_z = s_z
        self.t_xy = t_xy
        self.t_xz = t_xz
        self.t_yz = t_yz
        self.unit = unit

        # Calculate all derived values
        self._calculate()

    def _calculate(self):
        """Calculate stress invariants and principal stresses"""
        # Stress Invariants
        # Reference: https://en.wikiversity.org/wiki/Introduction_to_Elasticity/Principal_stresses
        self.I_1 = self.s_x + self.s_y + self.s_z
        self.I_2 = (self.s_x * self.s_y + self.s_y * self.s_z + self.s_z * self.s_x
                    - self.t_xy**2 - self.t_xz**2 - self.t_yz**2)
        self.I_3 = (self.s_x * self.s_y * self.s_z
                    - self.s_x * self.t_yz**2
                    - self.s_y * self.t_xz**2
                    - self.s_z * self.t_xy**2
                    + 2 * self.t_xy * self.t_xz * self.t_yz)

        # Calculate phi angle for principal stresses
        numerator = 2 * self.I_1**3 - 9 * self.I_1 * self.I_2 + 27 * self.I_3
        denominator = 2 * (self.I_1**2 - 3 * self.I_2)**(3/2)

        # Handle edge cases for acos
        cos_arg = numerator / (denominator + 1e-10)
        cos_arg = max(-1.0, min(1.0, cos_arg))  # Clamp to [-1, 1]

        phi = (1/3) * math.acos(cos_arg)

        # Principal Stresses
        self.s_1 = (self.I_1/3) + (2/3) * math.sqrt(self.I_1**2 - 3*self.I_2) * math.cos(phi)
        self.s_2 = (self.I_1/3) + (2/3) * math.sqrt(self.I_1**2 - 3*self.I_2) * math.cos(phi - 2*math.pi/3)
        self.s_3 = (self.I_1/3) + (2/3) * math.sqrt(self.I_1**2 - 3*self.I_2) * math.cos(phi - 4*math.pi/3)

        # Circle centers and radii
        self.C_1 = 0.5 * (self.s_1 + self.s_2)
        self.C_2 = 0.5 * (self.s_1 + self.s_3)
        self.C_3 = 0.5 * (self.s_2 + self.s_3)

        self.R_1 = abs(0.5 * (self.s_1 - self.s_2))
        self.R_2 = abs(0.5 * (self.s_1 - self.s_3))
        self.R_3 = abs(0.5 * (self.s_2 - self.s_3))

        # Von Mises stress
        self.s_vm = math.sqrt(0.5 * ((self.s_1 - self.s_3)**2 +
                                      (self.s_2 - self.s_3)**2 +
                                      (self.s_1 - self.s_2)**2))

        # Tresca (maximum shear stress)
        self.tau_max = self.R_2

    def generate_svg(self, filename="mohr_3d.svg", width=800, height=600):
        """
        Generate SVG visualization of Mohr's circle

        Args:
            filename: Output SVG file name
            width: SVG canvas width
            height: SVG canvas height
        """
        # Calculate bounds with padding
        padding = 50
        x_min = self.s_3 - padding
        x_max = self.s_1 + padding
        y_min = -self.R_2 - padding
        y_max = self.R_2 + padding

        # SVG coordinate transform parameters
        margin = 80
        plot_width = width - 2 * margin
        plot_height = height - 2 * margin

        def transform_x(x):
            """Transform data x-coordinate to SVG coordinate"""
            return margin + (x - x_min) / (x_max - x_min) * plot_width

        def transform_y(y):
            """Transform data y-coordinate to SVG coordinate (inverted)"""
            return margin + (y_max - y) / (y_max - y_min) * plot_height

        # Start building SVG
        svg_lines = [
            f'<?xml version="1.0" encoding="UTF-8"?>',
            f'<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg">',
            f'  <rect width="{width}" height="{height}" fill="white"/>',
            '',
            '  <!-- Mohr\'s Circle in 3D -->',
            ''
        ]

        # Draw axes
        origin_x = transform_x(0)
        origin_y = transform_y(0)

        # Y-axis
        svg_lines.append(f'  <!-- Y-axis -->')
        svg_lines.append(f'  <line x1="{origin_x}" y1="{transform_y(y_max)}" x2="{origin_x}" y2="{transform_y(y_min)}" stroke="black" stroke-width="2"/>')

        # X-axis
        svg_lines.append(f'  <!-- X-axis -->')
        svg_lines.append(f'  <line x1="{transform_x(x_min)}" y1="{origin_y}" x2="{transform_x(x_max)}" y2="{origin_y}" stroke="black" stroke-width="2"/>')

        # Draw axis ticks and labels
        unit_step = 100

        # X-axis ticks
        x_tick = math.floor(x_min / unit_step) * unit_step
        while x_tick <= x_max:
            x_pos = transform_x(x_tick)
            svg_lines.append(f'  <line x1="{x_pos}" y1="{origin_y-5}" x2="{x_pos}" y2="{origin_y+5}" stroke="black" stroke-width="1"/>')
            svg_lines.append(f'  <text x="{x_pos}" y="{origin_y+20}" text-anchor="middle" font-size="12" fill="black">{int(x_tick)}</text>')
            x_tick += unit_step

        # Y-axis ticks
        y_tick = math.floor(y_min / unit_step) * unit_step
        while y_tick <= y_max:
            if abs(y_tick) > 0.1:  # Skip zero
                y_pos = transform_y(y_tick)
                svg_lines.append(f'  <line x1="{origin_x-5}" y1="{y_pos}" x2="{origin_x+5}" y2="{y_pos}" stroke="black" stroke-width="1"/>')
                svg_lines.append(f'  <text x="{origin_x-10}" y="{y_pos+5}" text-anchor="end" font-size="12" fill="black">{int(y_tick)}</text>')
            y_tick += unit_step

        # Axis labels
        svg_lines.append(f'  <text x="{width/2}" y="{height-10}" text-anchor="middle" font-size="16" font-style="italic" fill="black">σ [{self.unit}]</text>')
        svg_lines.append(f'  <text x="20" y="{height/2}" text-anchor="middle" font-size="16" font-style="italic" fill="black" transform="rotate(-90 20 {height/2})">τ [{self.unit}]</text>')

        # Draw circles
        svg_lines.append('')
        svg_lines.append('  <!-- Mohr Circles -->')

        # Circle 1 (red)
        cx1 = transform_x(self.C_1)
        cy1 = transform_y(0)
        r1 = abs(transform_x(self.C_1 + self.R_1) - transform_x(self.C_1))
        svg_lines.append(f'  <circle cx="{cx1}" cy="{cy1}" r="{r1}" fill="none" stroke="red" stroke-width="2"/>')
        svg_lines.append(f'  <line x1="{cx1}" y1="{transform_y(self.R_1)}" x2="{cx1}" y2="{transform_y(-self.R_1)}" stroke="red" stroke-width="1"/>')

        # Circle 2 (blue)
        cx2 = transform_x(self.C_2)
        cy2 = transform_y(0)
        r2 = abs(transform_x(self.C_2 + self.R_2) - transform_x(self.C_2))
        svg_lines.append(f'  <circle cx="{cx2}" cy="{cy2}" r="{r2}" fill="none" stroke="blue" stroke-width="2"/>')
        svg_lines.append(f'  <line x1="{cx2}" y1="{transform_y(self.R_2)}" x2="{cx2}" y2="{transform_y(-self.R_2)}" stroke="blue" stroke-width="1"/>')

        # Circle 3 (green)
        cx3 = transform_x(self.C_3)
        cy3 = transform_y(0)
        r3 = abs(transform_x(self.C_3 + self.R_3) - transform_x(self.C_3))
        svg_lines.append(f'  <circle cx="{cx3}" cy="{cy3}" r="{r3}" fill="none" stroke="green" stroke-width="2"/>')
        svg_lines.append(f'  <line x1="{cx3}" y1="{transform_y(self.R_3)}" x2="{cx3}" y2="{transform_y(-self.R_3)}" stroke="green" stroke-width="1"/>')

        # Principal stress labels
        svg_lines.append('')
        svg_lines.append('  <!-- Principal Stress Labels -->')

        self._add_label(svg_lines, f"σ₁={self.s_1:.2f}", transform_x(self.s_1), origin_y - 20, "middle")
        self._add_label(svg_lines, f"σ₂={self.s_2:.2f}", transform_x(self.s_2), origin_y - 20, "middle")
        self._add_label(svg_lines, f"σ₃={self.s_3:.2f}", transform_x(self.s_3), origin_y - 20, "middle")

        # Circle center and radius labels
        svg_lines.append('')
        svg_lines.append('  <!-- Circle Parameters -->')

        self._add_label(svg_lines, f"({self.C_1:.2f}, {self.R_1:.2f})",
                       transform_x(self.C_1), transform_y(self.R_1) - 10, "middle", "red")
        self._add_label(svg_lines, f"({self.C_2:.2f}, {self.R_2:.2f})",
                       transform_x(self.C_2), transform_y(self.R_2) - 10, "middle", "blue")
        self._add_label(svg_lines, f"({self.C_3:.2f}, {self.R_3:.2f})",
                       transform_x(self.C_3), transform_y(self.R_3) - 10, "middle", "green")

        # Von Mises and Tresca labels
        svg_lines.append('')
        svg_lines.append('  <!-- Stress Metrics -->')

        info_x = transform_x(self.s_3)
        info_y = transform_y(self.R_2)

        self._add_label(svg_lines, f"σᵥₘ (Von Mises) = {self.s_vm:.2f}",
                       info_x, info_y - 30, "start", "black", 14)
        self._add_label(svg_lines, f"τₘₐₓ (Tresca) = {self.tau_max:.2f}",
                       info_x, info_y - 10, "start", "black", 14)

        # Close SVG
        svg_lines.append('</svg>')

        # Write to file
        with open(filename, 'w') as f:
            f.write('\n'.join(svg_lines))

        print(f"SVG generated: {filename}")

    def _add_label(self, svg_lines, text, x, y, anchor="middle", color="black", size=12):
        """Helper to add text label with background"""
        # Background rectangle
        text_width = len(text) * size * 0.6
        svg_lines.append(f'  <rect x="{x - text_width/2}" y="{y - size}" width="{text_width}" height="{size + 4}" fill="white" fill-opacity="0.8" stroke="none"/>')
        # Text
        svg_lines.append(f'  <text x="{x}" y="{y}" text-anchor="{anchor}" font-size="{size}" fill="{color}">{text}</text>')

    def print_summary(self):
        """Print calculation summary"""
        print("=" * 60)
        print("MOHR'S CIRCLE IN 3D - CALCULATION SUMMARY")
        print("=" * 60)
        print(f"\nInput Stresses ({self.unit}):")
        print(f"  σₓ = {self.s_x:10.2f}")
        print(f"  σᵧ = {self.s_y:10.2f}")
        print(f"  σᵤ = {self.s_z:10.2f}")
        print(f"  τₓᵧ = {self.t_xy:10.2f}")
        print(f"  τₓᵤ = {self.t_xz:10.2f}")
        print(f"  τᵧᵤ = {self.t_yz:10.2f}")

        print(f"\nStress Invariants:")
        print(f"  I₁ = {self.I_1:10.2f}")
        print(f"  I₂ = {self.I_2:10.2f}")
        print(f"  I₃ = {self.I_3:10.2f}")

        print(f"\nPrincipal Stresses ({self.unit}):")
        print(f"  σ₁ = {self.s_1:10.2f}")
        print(f"  σ₂ = {self.s_2:10.2f}")
        print(f"  σ₃ = {self.s_3:10.2f}")

        print(f"\nMohr Circle Parameters:")
        print(f"  Circle 1: Center = {self.C_1:10.2f}, Radius = {self.R_1:10.2f}")
        print(f"  Circle 2: Center = {self.C_2:10.2f}, Radius = {self.R_2:10.2f}")
        print(f"  Circle 3: Center = {self.C_3:10.2f}, Radius = {self.R_3:10.2f}")

        print(f"\nStress Metrics ({self.unit}):")
        print(f"  σᵥₘ (Von Mises)  = {self.s_vm:10.2f}")
        print(f"  τₘₐₓ (Tresca)    = {self.tau_max:10.2f}")
        print("=" * 60)


def main():
    """Main function with example stress values"""

    # Example 1: Default values from Mohr_3D.asy
    print("\nExample 1: Compression-dominated stress state")
    mohr1 = MohrCircle3D(
        s_x=-100,
        s_y=-200,
        s_z=-150,
        t_xy=50,
        t_xz=100,
        t_yz=150,
        unit="MPa"
    )
    mohr1.print_summary()
    mohr1.generate_svg("mohr_3d_example1.svg")

    print("\n")

    # Example 2: Alternative values from commented section
    print("\nExample 2: Mixed tension-compression stress state")
    mohr2 = MohrCircle3D(
        s_x=120,
        s_y=-85,
        s_z=0.0001,
        t_xy=-80,
        t_xz=1,
        t_yz=1,
        unit="MPa"
    )
    mohr2.print_summary()
    mohr2.generate_svg("mohr_3d_example2.svg")

    print("\n" + "=" * 60)
    print("SVG files generated successfully!")
    print("=" * 60)


if __name__ == "__main__":
    main()