Pure python linspace / logspace implementation

linspace.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# author: Panagiotis Mavrogiorgos
# email: gmail, pmav99

from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import absolute_import

import pytest


def _logspace(start, stop, num):
    num = int(num)
    start = start * 1.
    stop = stop * 1.

    if num <=1:
        raise ValueError("num must be greater than 1, not %d" % num)

    ratio = stop ** (1 / (num - 1))
    print(ratio)
    for i in range(num):
        if i == num - 1:
            yield stop
            return
        yield ratio ** i


def logspace(start, stop, num=10):
    diff = stop - start
    for i, value in enumerate(_logspace(start, stop, num)):
        if i == 0:
            yield start
        elif i == num:
            yield stop
        else:
            yield start + value * diff / stop


def linspace(start, stop, num=50, endpoint=True):
    num = int(num)
    start = start * 1.
    stop = stop * 1.

    if num == 1:
        yield stop
        return
    if endpoint:
        step = (stop - start) / (num - 1)
    else:
        step = (stop - start) / num

    for i in range(num):
        yield start + step * i


def test_linspace():
    assert list(linspace(0, 10, 3, True)) == [0, 5, 10]
    assert list(linspace(0, 10, 5, True)) == [0, 2.5, 5, 7.5, 10]
    assert list(linspace(0, 10, 11, True)) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

    assert list(linspace(1, 10, 3, False)) == [1, 4, 7]
    assert list(linspace(0, 10, 5, False)) == [0, 2, 4, 6, 8]
    assert list(linspace(0, 10, 10, False)) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]