How to use the omas.omas_physics.cocos_transform function in omas

def pick_direction():
        """Guesses the direction for the arrow (from injector toward machine) in case you don't know"""
        dr = ods['equilibrium']['time_slice'][0]['global_quantities']['magnetic_axis']['r'] - r
        dz = ods['equilibrium']['time_slice'][0]['global_quantities']['magnetic_axis']['z'] - z
        theta = numpy.arctan2(dz, -dr)
        if snap_to > 0:
            theta = snap_to * round(theta / snap_to)
        return theta

    if (r2 is not None) and (z2 is not None):
        direction = numpy.arctan2(z2 - z, r - r2)
    elif direction is None:
        direction = pick_direction()
    else:
        direction = cocos_transform(ods.cocos, 11)['BP'] * direction

    if ax is None:
        ax = pyplot.gca()

    shaft_len = 3.5 * (1 + pad) / 2.

    da = numpy.pi / 10  # Angular half width of the arrow head
    x0 = numpy.cos(-direction) * pad
    y0 = numpy.sin(-direction) * pad
    head_mark = [
        (x0, y0),
        (x0 + numpy.cos(-direction + da), y0 + numpy.sin(-direction + da)),
        (x0 + numpy.cos(-direction), y0 + numpy.sin(-direction)),
        (x0 + shaft_len * numpy.cos(-direction), y0 + shaft_len * numpy.sin(-direction)),
        (x0 + numpy.cos(-direction), y0 + numpy.sin(-direction)),
        (x0 + numpy.cos(-direction - da), y0 + numpy.sin(-direction - da)),