<commit_before>'Rayleigh-Ritz solution to fully-coupled uvw-phi equations.'

import numpy as np
from numpy import pi
from scipy.optimize import minimize


class ModeMatrix:
    """Solve coupled lateral, radial, and torsional deflections."""

    def B_theta(self, theta=[0.], comps=[0, 1, 2, 3]):
        'Matrix to transform mode coefficients to vector components.'

        # Convert scalar values to arrays
        theta = np.atleast_1d(theta)
        comps = np.atleast_1d(comps)

        B = np.zeros((len(theta)*len(comps), 4 + 8*self.n_modes))

        # For each angle theta
        for it in range(len(theta)):

            # Zero mode
            for ic, c in enumerate(comps):
                B[len(comps)*it + ic, c] = 1.

            # Higher modes
            for n in range(1, self.n_modes + 1):
                cos_ni = np.cos(n*theta[it])
                sin_ni = np.sin(n*theta[it])

                for ic, c in enumerate(comps):
                    B[len(comps)*it + ic, 4 + 8*(n-1) + 2*c] = cos_ni
                    B[len(comps)*it + ic, 4 + 8*(n-1) + 2*c+1] = sin_ni

        return B

    def K_rim_matl(self, r0=True):
        'Elastic portion of K_rim.'

        w = self.wheel

        R = w.rim.radius                        # rim radius
        EA = w.rim.young_mod * w.rim.area       # axial stiffness
        EI_rad = w.rim.young_mod * w.rim.I_rad  # radial bending
        EI_lat = w.rim.young_mod * w.rim.I_lat  # lateral bending
        EIw = w.rim.young_mod * w.rim.I_warp    # warping constant
        GJ = w.rim.shear_mod * w.rim.J_tor      # torsion constant

        y0 = 0.  # shear-center offset
        if 'y_0' in w.rim.sec_params:
            y0 = w.rim.sec_params['y_0']

        r02 = 0.
        if r0:
            r02 = EI_rad/EA + EI_lat/EA + y0**2

        K_rim_matl = np.zeros((4 + self.n_modes*8, 4 + self.n_modes*8))

        # zero mode
        K_rim_matl[1, 1] = 2*pi*EA/R
        K_rim_matl[3, 3] = 2*pi*EI_lat/R

        # higher modes
        for n in range(1, self.n_modes + 1):
            i0 = 4 + (n-1)*8

            # k_vv
            K_rim_matl[i0+2, i0+2] = EI_rad*pi/R**3*n**4 + EA*pi/R*(1 + y0/R*n**2)**2
            K_rim_matl[i0+3, i0+3] = K_rim_matl[i0+2, i0+2]

            # k_ww
            K_rim_matl[i0+4, i0+4] = EI_rad*pi/R**3*n**2 + EA*pi/R*n**2*(1 + y0/R)**2
            K_rim_matl[i0+5, i0+5] = K_rim_matl[i0+4, i0+4]

            # k_vw
            K_rim_matl[i0+2, i0+5] = -EI_rad*pi/R**3*n**3 -\
                EA*pi*n/R*(1 + y0/R*(1 + n**2) + y0**2/R**2*n**2)
            K_rim_matl[i0+5, i0+2] = K_rim_matl[i0+2, i0+5]
            K_rim_matl[i0+3, i0+4] = -K_rim_matl[i0+2, i0+5]
            K_rim_matl[i0+4, i0+3] = -K_rim_matl[i0+2, i0+5]

            # k_uu
            K_rim_matl[i0+0, i0+0] = (EI_lat*pi/R**3*n**4 + EIw*pi/R**5*n**4 +
                                      GJ*pi/R**3*n**2)
            K_rim_matl[i0+1, i0+1] = K_rim_matl[i0+0, i0+0]

            # k_up
            K_rim_matl[i0+0, i0+6] = -(EI_lat*pi/R**2*n**2 + EIw*pi/R**4*n**4 +
                                  GJ*pi/R**2*n**2)
            K_rim_matl[i0+6, i0+0] = K_rim_matl[i0+0, i0+6]
            K_rim_matl[i0+1, i0+7] = K_rim_matl[i0+0, i0+6]
            K_rim_matl[i0+7, i0+1] = K_rim_matl[i0+0, i0+6]

            # k_pp
            K_rim_matl[i0+6, i0+6] = (EI_lat*pi/R + EIw*pi/R**3*n**4 + GJ*pi/R*n**2)
            K_rim_matl[i0+7, i0+7] = K_rim_matl[i0+6, i0+6]

        return K_rim_matl

    def K_rim_geom(self, r0=True):
        'Tension-dependent portion of K_rim, such that K_rim = K_rim_matl - T_avg*K_rim_geom'

        w = self.wheel

        R = w.rim.radius

        y0 = 0.  # shear-center offset
        if 'y_0' in w.rim.sec_params:
            y0 = w.rim.sec_params['y_0']

        r02 = 0.
        if r0:
            r02 = w.rim.I_rad/w.rim.area + w.rim.I_lat/w.rim.area + y0**2

        K_rim_geom = np.zeros((4 + self.n_modes*8, 4 + self.n_modes*8))

        # zero mode
        K_rim_geom[3, 3] = -2.*pi*R*y0

        # higher modes
        for n in range(1, self.n_modes + 1):
            i0 = 4 + (n-1)*8

            # k_uu
            K_rim_geom[i0+0, i0+0] = pi*n**2*(1. + r02/R**2)
            K_rim_geom[i0+1, i0+1] = K_rim_geom[i0+0, i0+0]

            # k_up
            K_rim_geom[i0+0, i0+6] = pi*n**2*(y0 - r02/R*n**2)
            K_rim_geom[i0+6, i0+0] = K_rim_geom[i0+0, i0+6]<commit_msg>Fix minor typo in K_rim<commit_after>'Rayleigh-Ritz solution to fully-coupled uvw-phi equations.'

import numpy as np
from numpy import pi
from scipy.optimize import minimize


class ModeMatrix:
    """Solve coupled lateral, radial, and torsional deflections."""

    def B_theta(self, theta=[0.], comps=[0, 1, 2, 3]):
        'Matrix to transform mode coefficients to vector components.'

        # Convert scalar values to arrays
        theta = np.atleast_1d(theta)
        comps = np.atleast_1d(comps)

        B = np.zeros((len(theta)*len(comps), 4 + 8*self.n_modes))

        # For each angle theta
        for it in range(len(theta)):

            # Zero mode
            for ic, c in enumerate(comps):
                B[len(comps)*it + ic, c] = 1.

            # Higher modes
            for n in range(1, self.n_modes + 1):
                cos_ni = np.cos(n*theta[it])
                sin_ni = np.sin(n*theta[it])

                for ic, c in enumerate(comps):
                    B[len(comps)*it + ic, 4 + 8*(n-1) + 2*c] = cos_ni
                    B[len(comps)*it + ic, 4 + 8*(n-1) + 2*c+1] = sin_ni

        return B

    def K_rim_matl(self, r0=True):
        'Elastic portion of K_rim.'

        w = self.wheel

        R = w.rim.radius                        # rim radius
        EA = w.rim.young_mod * w.rim.area       # axial stiffness
        EI_rad = w.rim.young_mod * w.rim.I_rad  # radial bending
        EI_lat = w.rim.young_mod * w.rim.I_lat  # lateral bending
        EIw = w.rim.young_mod * w.rim.I_warp    # warping constant
        GJ = w.rim.shear_mod * w.rim.J_tor      # torsion constant

        y0 = 0.  # shear-center offset
        if 'y_0' in w.rim.sec_params:
            y0 = w.rim.sec_params['y_0']

        r02 = 0.
        if r0:
            r02 = EI_rad/EA + EI_lat/EA + y0**2

        K_rim_matl = np.zeros((4 + self.n_modes*8, 4 + self.n_modes*8))

        # zero mode
        K_rim_matl[1, 1] = 2*pi*EA/R
        K_rim_matl[3, 3] = 2*pi*EI_lat/R

        # higher modes
        for n in range(1, self.n_modes + 1):
            i0 = 4 + (n-1)*8

            # k_vv
            K_rim_matl[i0+2, i0+2] = EI_rad*pi/R**3*n**4 + EA*pi/R*(1 + y0/R*n**2)**2
            K_rim_matl[i0+3, i0+3] = K_rim_matl[i0+2, i0+2]

            # k_ww
            K_rim_matl[i0+4, i0+4] = EI_rad*pi/R**3*n**2 + EA*pi/R*n**2*(1 + y0/R)**2
            K_rim_matl[i0+5, i0+5] = K_rim_matl[i0+4, i0+4]

            # k_vw
            K_rim_matl[i0+2, i0+5] = -EI_rad*pi/R**3*n**3 -\
                EA*pi*n/R*(1 + y0/R*(1 + n**2) + y0**2/R**2*n**2)
            K_rim_matl[i0+5, i0+2] = K_rim_matl[i0+2, i0+5]
            K_rim_matl[i0+3, i0+4] = -K_rim_matl[i0+2, i0+5]
            K_rim_matl[i0+4, i0+3] = -K_rim_matl[i0+2, i0+5]

            # k_uu
            K_rim_matl[i0+0, i0+0] = (EI_lat*pi/R**3*n**4 + EIw*pi/R**5*n**4 +
                                      GJ*pi/R**3*n**2)
            K_rim_matl[i0+1, i0+1] = K_rim_matl[i0+0, i0+0]

            # k_up
            K_rim_matl[i0+0, i0+6] = -(EI_lat*pi/R**2*n**2 + EIw*pi/R**4*n**4 +
                                  GJ*pi/R**2*n**2)
            K_rim_matl[i0+6, i0+0] = K_rim_matl[i0+0, i0+6]
            K_rim_matl[i0+1, i0+7] = K_rim_matl[i0+0, i0+6]
            K_rim_matl[i0+7, i0+1] = K_rim_matl[i0+0, i0+6]

            # k_pp
            K_rim_matl[i0+6, i0+6] = (EI_lat*pi/R + EIw*pi/R**3*n**4 + GJ*pi/R*n**2)
            K_rim_matl[i0+7, i0+7] = K_rim_matl[i0+6, i0+6]

        return K_rim_matl

    def K_rim_geom(self, r0=True):
        'Tension-dependent portion of K_rim, such that K_rim = K_rim_matl - T_avg*K_rim_geom'

        w = self.wheel

        R = w.rim.radius

        y0 = 0.  # shear-center offset
        if 'y_0' in w.rim.sec_params:
            y0 = w.rim.sec_params['y_0']

        r02 = 0.
        if r0:
            r02 = w.rim.I_rad/w.rim.area + w.rim.I_lat/w.rim.area + y0**2

        K_rim_geom = np.zeros((4 + self.n_modes*8, 4 + self.n_modes*8))

        # zero mode
        K_rim_geom[3, 3] = -2.*pi*R*y0

        # higher modes
        for n in range(1, self.n_modes + 1):
            i0 = 4 + (n-1)*8

            # k_uu
            K_rim_geom[i0+0, i0+0] = pi*n**2*(1. + r02/R**2)
            K_rim_geom[i0+1, i0+1] = K_rim_geom[i0+0, i0+0]

            # k_up
            K_rim_geom[i0+0, i0+6] = pi*n**2*(y0 - r02/R)
            K_rim_geom[i0+6, i0+0] = K_rim_geom[i0+0, i0+6]
{'fraction_non_alphanumeric': 0.1575957277582093, 'fraction_numerical': 0.05760709010339734, 'mean_word_length': 1.7024869511820695, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1852713', 'n_tokens_mistral': 4115, 'n_tokens_neox': 3876, 'n_words': 679}